From 526edcb50eee73cabe6f8d5bbed3921be08ff20c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 9 Apr 2026 12:09:52 -0300 Subject: [PATCH] feat: update training workflow and repository management - Replaced synchronous MinIO repository calls with asynchronous counterparts in the Training class for improved performance. - Enhanced logging throughout the training process to provide better insights into model metadata loading, parameter validation, and training execution. - Updated the train_test_split function to enforce DataFrame input type, ensuring consistency in data handling. - Removed the deprecated model_repository.py file to streamline the codebase. - Adjusted cleanup schedule logic to improve error handling and logging during schedule reconciliation. - Updated tests to reflect changes in the training workflow and repository interactions. --- model_manager/activities/training.py | 40 +- model_manager/schedules/cleanup_schedule.py | 5 +- .../repository/data_manager_repository.py | 9 +- .../utils/repository/model_repository.py | 471 ------------------ model_manager/worker/worker.py | 5 +- model_manager/workflows/train_model.py | 2 + requirements-dev.txt | 1 + scripts/run_training_test.py | 22 +- tests/activities/test_training.py | 12 +- tests/schedules/test_cleanup_schedule.py | 23 + tests/sientia/test_exceptions.py | 9 + .../test_data_manager_repository.py | 8 +- tests/worker/test_prepare_worker.py | 6 + tests/workflows/test_train_model.py | 4 +- 14 files changed, 114 insertions(+), 503 deletions(-) delete mode 100644 model_manager/utils/repository/model_repository.py create mode 100644 tests/sientia/test_exceptions.py diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index e33e4a3..506a4b5 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -18,7 +18,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring - from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.repository.minio_repository_sync import MinioRepository from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository from sientia_model.model_repository.plugin_store import PluginStore @@ -78,14 +78,18 @@ class Training(SientiaMonitoring): """ metadata = input_data.get('metadata', {}) + self.info(f'Loading model metadata for {input_data}', metadata) + try: train_params = TrainModelParams.from_dict(input_data) model_metadata = self.plugin_store.get_model_index( - model_name=train_params.model_name, + model_type=train_params.model_type, metadata=metadata, ) train_params.model_metadata = model_metadata + self.info(f'Model metadata loaded successfully for {input_data}', metadata) + self.debug(f'Model metadata: {model_metadata}', metadata) return train_params.to_dict() except Exception as exc: trace = traceback.format_exc() @@ -120,6 +124,9 @@ class Training(SientiaMonitoring): Exception: If validation fails (after sending notification) """ metadata = input_data.get('metadata', {}) + + self.info(f'Validating training parameters for {input_data}', metadata) + try: train_params = TrainModelParams.from_dict(input_data) @@ -132,6 +139,10 @@ class Training(SientiaMonitoring): metadata, ) + self.debug( + f'Training parameters validated successfully: {train_params.to_dict()}', metadata + ) + return train_params.to_dict() except Exception as e: error_msg = f'Error validating training parameters: {str(e)}' @@ -173,9 +184,15 @@ class Training(SientiaMonitoring): metadata = input_data.get('metadata') train_params = TrainModelParams.from_dict(input_data['train_params']) + self.info('Starting train_model process', metadata) + try: # Download training file bytes from MinIO - train_bytes = self.minio_repository.download_file_sync( + + self.info( + f'Downloading training file from MinIO for {train_params.file_name}', metadata + ) + train_bytes = self.minio_repository.download_file( object_name=train_params.file_name, bucket=train_params.bucket_name, metadata=metadata, @@ -185,12 +202,14 @@ class Training(SientiaMonitoring): val_bytes: bytes | None = None validation_name = train_params.val_file_name if validation_name is not None: - val_bytes = self.minio_repository.download_file_sync( + self.info(f'Downloading validation file from MinIO for {validation_name}', metadata) + val_bytes = self.minio_repository.download_file( object_name=validation_name, bucket=train_params.bucket_name, metadata=metadata, ) + self.info(f'Preparing training data for {train_params.file_name}', metadata) train_result = self.data_manager_repository.prepare_training_data( train_file_bytes=train_bytes, validation_file_bytes=val_bytes, @@ -198,8 +217,9 @@ class Training(SientiaMonitoring): metadata=metadata, ) + self.info(f'Getting model wrapper for {train_params.model_type}', metadata) wrapper = self.plugin_store.get_model( - model_name=train_params.model_name, + model_type=train_params.model_type, force_download=False, opt_params=train_params.opt_params or {}, model_kwargs=train_params.model_kwargs or {}, @@ -207,6 +227,7 @@ class Training(SientiaMonitoring): metadata=metadata, ) + self.info(f'Training model for {train_params.model_type}', metadata) train_data = train_result.train_data val_data = train_result.val_data @@ -216,6 +237,10 @@ class Training(SientiaMonitoring): target=train_params.target_variable, ) + self.info( + f'Generating predictions using the trained wrapper for {train_params.model_type}', + metadata, + ) # Generate predictions using the trained wrapper transformed_train, _ = wrapper.transform(train_data) transformed_val, _ = wrapper.transform(val_data) @@ -229,11 +254,13 @@ class Training(SientiaMonitoring): train_result.y_train_pred = y_train_pred_df train_result.y_pred = y_val_pred_df + self.info(f'Computing regression metrics for {train_params.model_type}', metadata) train_result = self.data_manager_repository.compute_regression_metrics( train_result, wrapper, ) + self.info(f'Starting MLflow run for {train_params.model_type}', metadata) with self.mlflow_repository.start_run( model_name=train_params.model_name, run_name=None, @@ -273,6 +300,7 @@ class Training(SientiaMonitoring): wrapper: Any, metadata: dict[str, Any] | None, ) -> None: + self.info(f'Generating report for {train_params.model_type}', metadata) train_result = self.data_manager_repository.generate_report( train_result, metadata=metadata, @@ -285,7 +313,9 @@ class Training(SientiaMonitoring): ): raise ValueError('Report path, train data path, or test data path is not set') + self.info(f'Storing model for {train_params.model_type}', metadata) wrapper.store_model(name=train_params.model_name) + self.info(f'Logging artifacts for {train_params.model_type}', metadata) mlflow.log_artifact(train_result.report_path) mlflow.log_artifact(train_result.train_data_path) mlflow.log_artifact(train_result.test_data_path) diff --git a/model_manager/schedules/cleanup_schedule.py b/model_manager/schedules/cleanup_schedule.py index c41b4f3..336fda5 100644 --- a/model_manager/schedules/cleanup_schedule.py +++ b/model_manager/schedules/cleanup_schedule.py @@ -116,7 +116,6 @@ async def create_cleanup_schedule( runtime = (os.getenv('RUNTIME') or 'single').strip() cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single') schedule_id = build_cleanup_schedule_id(runtime) - created = False updated = False if await schedule_exists(client, schedule_id, logger, metadata): @@ -147,15 +146,13 @@ async def create_cleanup_schedule( ), ), ) - created = not updated - if updated: logger.custom_info( f"Schedule '{schedule_id}' reconciled successfully. " f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', metadata, ) - elif created: + else: logger.custom_info( f"Schedule '{schedule_id}' created successfully. " f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py index 33a01d0..63e6a57 100644 --- a/model_manager/utils/repository/data_manager_repository.py +++ b/model_manager/utils/repository/data_manager_repository.py @@ -34,7 +34,7 @@ from model_manager.utils.models.train_model_result import TrainModelResult def train_test_split( - data: pd.DataFrame | pd.Series, + data: pd.DataFrame, train_size: float, random_state: int | None = None, shuffle: bool = True, @@ -57,11 +57,8 @@ def train_test_split( train_indices = indices[:n_train] test_indices = indices[n_train:] - # 5. Retornar os dados fatiados (funciona para DataFrame ou Series) - if isinstance(data, (pd.DataFrame, pd.Series)): - return data.iloc[train_indices], data.iloc[test_indices] - - return data[train_indices], data[test_indices] + # 5. Retornar os dados fatiados + return data.iloc[train_indices], data.iloc[test_indices] def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame: diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py deleted file mode 100644 index dcd1ba8..0000000 --- a/model_manager/utils/repository/model_repository.py +++ /dev/null @@ -1,471 +0,0 @@ -""" -MLFlow Repository - -This module contains the MLFlowRepository class, which is responsible for -handling model training artifacts and MLFlow operations for the Model Manager system. - -It includes methods for generating training reports, managing artifacts, -and logging model runs to MLFlow. - -""" - -import json -import os -import shutil -import warnings -from datetime import datetime -from os import makedirs, path - -import numpy as np -import pandas as pd -from sientia_do.observability.logger import Logger - -from model_manager.sientia.model_serving import ModelServing # type: ignore[import-untyped] -from model_manager.sientia.reports import Reports # type: ignore[import-untyped] -from model_manager.utils.models.train_model_result import TrainModelResult - -# Suppress sklearn FutureWarning about 'squared' deprecation without changing business logic -warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' is deprecated.*") - - -class ModelRepository: - def __init__(self, url, username, password, logger: Logger): - self.model_serving = ModelServing(tracking_uri=url, username=username, password=password) - self.logger = logger - self.logger.info(f'MLFlow client initialized at {url}') - - def save_model(self, train_result: TrainModelResult) -> TrainModelResult: - """ - Save a trained ML model and its artifacts to MLflow. - - This activity orchestrates the complete model saving pipeline: - 1. Generates the next run name for the experiment - 2. Creates and organizes artifacts (reports, data files) - 3. Logs model, parameters, metrics, and artifacts to MLflow - - Args: - input_data: Configuration for model saving operation - Required keys: - - metadata (dict): Workflow execution metadata - - train_result (TrainModelResult): Training result with model and metrics - - Returns: - TrainModelResult: Updated training result with run_name and artifacts - - Raises: - Exception: If model saving fails (after sending notification) - """ - experiment_name = train_result.params.experiment_name - train_result.run_name = self._get_next_run_name(experiment_name) - train_result = self._generate_artifacts(train_result) - self._save_run(train_result) - - self.logger.info( - f'Model saved successfully - experiment run id: {train_result.params.experiment_run_id}, ' - f'experiment name: {experiment_name}, ' - f'run name: {train_result.run_name}' - ) - - return train_result - - def cleanup_run_directory(self, run_dir: str) -> None: - """ - Clean up temporary run directory after model training. - - This activity deletes the temporary directory created during model training - and artifact generation. It implements idempotent cleanup to handle cases - where the directory may have already been deleted. - - Args: - run_dir (str): Path to the run directory to delete - """ - if not run_dir: - self.logger.info('No run directory specified, skipping cleanup') - return - - if os.path.exists(run_dir): - shutil.rmtree(run_dir) - self.logger.info(f'Run directory deleted successfully: {run_dir}') - else: - self.logger.info(f'Run directory already deleted: {run_dir}') - - def _get_next_run_name(self, experiment_name: str) -> str: - """ - Generates the next run name for a given experiment. - - Args: - experiment_name (str): The name of the experiment for which the next run name is being generated. - - Returns: - str: A unique run name in the format "-". - """ - runs = self.model_serving.search_runs_by_name( - experiment_names=[experiment_name], order_by=['start_time desc'] - ) - - next_run_number = len(runs) + 1 - return f'{experiment_name}-{next_run_number}' - - def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult: - """ - Generates and organizes artifacts related to the training process, such as reports and data files. - - Args: - data: The training model result containing the datasets, model, and parameters. - - Returns: - The updated result object with paths to the generated artifacts. - - Raises: - FileNotFoundError: If the reports directory or header.html file does not exist. - ValueError: If run_name is not set. - """ - # Validate that run_name is set - if not data.run_name: - error_msg = 'run_name must be set before generating artifacts' - self.logger.error(error_msg) - raise ValueError(error_msg) - - reference_data, current_data = self._init_artifacts_data(data) - base_path = self._get_reports_directory() - - # Validate that reports directory exists - if not path.exists(base_path): - error_msg = f'Reports directory does not exist: {base_path}' - self.logger.error(error_msg) - raise FileNotFoundError(error_msg) - - data.run_dir = self._create_run_directory(base_path, data.run_name) - header_file_path = path.join(base_path, 'header.html') - - # Validate that header.html exists - if not path.exists(header_file_path): - error_msg = f'Header file does not exist: {header_file_path}' - self.logger.error(error_msg) - raise FileNotFoundError(error_msg) - - self._setup_run_directory(data.run_dir, header_file_path) - return self._generate_report(reference_data, current_data, data) - - def _save_run(self, data: TrainModelResult): - """ - Logs the details of a machine learning run, including parameters, metrics, models, and artifacts, - to the Sientia tracking system. - - Args: - data: The training model result containing the datasets, model, parameters, - and evaluation metrics. - - Raises: - ValueError: If required metrics or artifacts are missing. - Exception: If MLflow logging fails for any reason. - """ - # Validate that required artifacts exist before attempting to log - if not data.report_path or not path.exists(data.report_path): - error_msg = f'Report file does not exist: {data.report_path}' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if not data.train_data_path or not path.exists(data.train_data_path): - error_msg = f'Training data file does not exist: {data.train_data_path}' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if not data.test_data_path or not path.exists(data.test_data_path): - error_msg = f'Test data file does not exist: {data.test_data_path}' - self.logger.error(error_msg) - raise ValueError(error_msg) - - # Validate that metrics are present - if data.mse_val is None or data.r2_val is None or data.mae_val is None: - error_msg = 'One or more metrics (MSE, R2, MAE) are None' - self.logger.error(error_msg) - raise ValueError(error_msg) - - # Prepare parameters - interval_strs = [ - (str(interval[0]), str(interval[1])) - for interval in (data.params.removed_intervals or []) - ] - - # Set experiment and create run - self.model_serving.set_experiment(data.params.experiment_name) - - with self.model_serving.save_experiment( - run_name=data.run_name, description=data.params.experiment_name - ): - # Log model parameters - self.model_serving.log_param('model_name', data.params.model_name) - self.model_serving.log_param( - 'models_params', - {'degree': data.params.degree, 'interaction_only': data.params.interaction_only}, - ) - self.model_serving.log_param('target_variable', data.params.target_variable) - self.model_serving.log_param('input_variables', data.params.variable_columns) - self.model_serving.log_param('nan_treatment', data.params.nan_treatment) - self.model_serving.log_param('lag_train', data.params.lag_train) - self.model_serving.log_param('lag_transform', data.params.lag_val) - static_threshold_value = None - if data.params.rem_static_win: - static_threshold_value = ( - data.params.static_threshold if data.params.static_threshold is not None else 1 - ) - self.model_serving.log_param('static_threshold', static_threshold_value) - self.model_serving.log_param('lower_limits', data.params.low_lim) - self.model_serving.log_param('upper_limits', data.params.upp_lim) - self.model_serving.log_param('scaler_name', data.params.scaler_name) - self.model_serving.log_param('scaler_params', data.scaler_dict) - self.model_serving.log_param('include_ar', data.params.include_ar) - self.model_serving.log_param('train_size', round(data.params.train_size / 100, 2)) - self.model_serving.log_param('test_size', round(1 - (data.params.train_size / 100), 2)) - self.model_serving.log_param('start_date', data.params.start_date) - self.model_serving.log_param('end_date', data.params.end_date) - self.model_serving.log_param('removed_intervals', interval_strs) - self.model_serving.log_param('retrain', False) - self.model_serving.log_param('support_filters', data.params.support_filters) - - # Log evaluation metrics - self.model_serving.log_metric('MSE', data.mse_val) - self.model_serving.log_metric('R2', data.r2_val) - self.model_serving.log_metric('MAE', data.mae_val) - - # Log models - self.model_serving.log_model(data.process_data, 'data_model') - self.model_serving.log_model(data.regr, 'prediction_model') - - # Log artifacts - self.model_serving.log_artifact(data.report_path) - self.model_serving.log_artifact(data.train_data_path) - self.model_serving.log_artifact(data.test_data_path) - - # Log equation artifact if available - if data.equation_path and path.exists(data.equation_path): - self.model_serving.log_artifact(data.equation_path) - - def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]: - """ - Prepares the reference and current datasets for artifact generation. - - Args: - data: The training model result containing the datasets and model. - - Returns: - tuple: A tuple containing: - - reference_data: The training dataset with predictions added. - - current_data: The testing dataset with predictions added. - - Raises: - ValueError: If training or test datasets are empty or invalid. - AttributeError: If required attributes are missing from the data object. - """ - # Validate that required DataFrames are not empty - # Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult - # so we only check if they are empty, not None - if data.x_train.empty: - error_msg = 'Training features (x_train) are empty' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if data.y_train.empty: - error_msg = 'Training target (y_train) is empty' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if data.x_test.empty: - error_msg = 'Test features (x_test) are empty' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if data.y_test.empty: - error_msg = 'Test target (y_test) is empty' - self.logger.error(error_msg) - raise ValueError(error_msg) - - # Validate that predictions exist (y_pred is optional, so check for None) - if data.y_pred is None: - error_msg = 'Test predictions (y_pred) are None' - self.logger.error(error_msg) - raise ValueError(error_msg) - - if data.y_train_pred is None: - error_msg = 'Training predictions (y_train_pred) are None' - self.logger.error(error_msg) - raise ValueError(error_msg) - - # Prepare reference data (training set) - reference_data = pd.concat([data.x_train, data.y_train], axis=1) - reference_data = reference_data.rename(columns={data.params.target_variable: 'target'}) - # Use pre-calculated predictions (calculated before denormalization to avoid overflow) - reference_data['prediction'] = data.y_train_pred - - # Prepare current data (test set) - current_data = pd.concat([data.x_test, data.y_test], axis=1) - current_data = current_data.rename(columns={data.params.target_variable: 'target'}) - current_data['prediction'] = data.y_pred - - return reference_data, current_data - - def _create_run_directory(self, base_path: str, run_name: str) -> str: - """ - Creates a directory inside the 'reports' folder with the run name and a timestamp. - - Uses microsecond precision in timestamp to minimize collision probability - in high-concurrency scenarios. - - Args: - base_path (str): The path to the 'reports' folder. - run_name (str): The name of the run. - - Returns: - str: The path to the created directory. - - Raises: - PermissionError: If there are insufficient permissions to create the directory. - OSError: If directory creation fails for any other reason. - """ - # Use microsecond precision to reduce collision probability - timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f') - run_dir = path.join(base_path, 'temp', f'{run_name}_{timestamp}') - - try: - makedirs(run_dir, exist_ok=True) - return run_dir - except PermissionError as e: - error_msg = f'Permission denied when creating directory: {run_dir}' - self.logger.error(error_msg) - raise PermissionError(error_msg) from e - except OSError as e: - error_msg = f'Failed to create directory {run_dir}: {str(e)}' - self.logger.error(error_msg) - raise OSError(error_msg) from e - - def _setup_run_directory(self, run_dir: str, header_file_path: str): - """ - Creates empty files and copies a header file into the specified run directory. - - Note: Lock removed as each run has its own unique directory, so no synchronization - is needed between different runs. File operations within the same directory are - atomic at the OS level. - - Args: - run_dir (str): The path to the run directory where the files will be created. - header_file_path (str): The path to the header.html file to be copied. - - Raises: - FileNotFoundError: If the header file does not exist. - PermissionError: If there are insufficient permissions to create files. - OSError: If file creation or copying fails for any other reason. - """ - empty_files = ['data_drift.html', 'data_quality.html', 'regression.html'] - - try: - # Create empty placeholder files - for file_name in empty_files: - file_path = path.join(run_dir, file_name) - with open(file_path, 'w'): - pass # Create empty file - - # Copy header file to run directory - header_dest = path.join(run_dir, 'header.html') - shutil.copy(header_file_path, header_dest) - except FileNotFoundError as e: - error_msg = f'Header file not found: {header_file_path}' - self.logger.error(error_msg) - raise FileNotFoundError(error_msg) from e - except PermissionError as e: - error_msg = f'Permission denied when setting up directory: {run_dir}' - self.logger.error(error_msg) - raise PermissionError(error_msg) from e - except OSError as e: - error_msg = f'Failed to setup run directory {run_dir}: {str(e)}' - self.logger.error(error_msg) - raise OSError(error_msg) from e - - def _generate_report( - self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult - ) -> TrainModelResult: - """ - Generates a comprehensive report summarizing data quality, data drift, and regression analysis. - - Args: - reference_data (pd.DataFrame): The training dataset with predictions added. - current_data (pd.DataFrame): The testing dataset with predictions added. - data: The training model result containing the datasets, model, and parameters. - - Returns: - The updated result object with paths to the generated report and data files. - - Raises: - ValueError: If data conversion to float64 fails or DataFrames are invalid. - PermissionError: If there are insufficient permissions to write files. - OSError: If file writing fails for any other reason. - """ - try: - # Convert data to float64 for report generation - # This may raise ValueError if data contains non-numeric values - reference_data_float = reference_data.astype(np.float64) - current_data_float = current_data.astype(np.float64) - - # Initialize report generator - report = Reports( - reference_data=reference_data_float, - current_data=current_data_float, - base_path=data.run_dir, - ) - - # Generate report sections - report.add_data_quality_section(columns=data.params.variable_columns + ['target']) - report.add_data_drift_section(columns=data.params.variable_columns + ['target']) - report.add_regression_section() - - # Validate that run_dir is set (should be set by _create_run_directory) - if not data.run_dir: - error_msg = 'run_dir is not set after directory creation' - self.logger.error(error_msg) - raise ValueError(error_msg) - - # Save HTML report - data.report_path = path.join(data.run_dir, 'report.html') - report.save_all_sections_html(data.report_path) - - # Save training data CSV - data.train_data_path = path.join(data.run_dir, 'train_data.csv') - reference_data.to_csv(data.train_data_path, index=False) - - # Save test data CSV - data.test_data_path = path.join(data.run_dir, 'test_data.csv') - current_data.to_csv(data.test_data_path, index=False) - - # Save equation as JSON - if data.equation is not None: - data.equation_path = path.join(data.run_dir, 'model_equation.json') - with open(data.equation_path, 'w', encoding='utf-8') as f: - json.dump(data.equation, f, indent=2, ensure_ascii=False) - - return data - except ValueError as e: - error_msg = f'Failed to convert data to float64 for report generation: {str(e)}' - self.logger.error(error_msg) - raise ValueError(error_msg) from e - except PermissionError as e: - error_msg = f'Permission denied when writing report files to: {data.run_dir}' - self.logger.error(error_msg) - raise PermissionError(error_msg) from e - except OSError as e: - error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}' - self.logger.error(error_msg) - raise OSError(error_msg) from e - - def _get_reports_directory(self) -> str: - """ - Get the absolute path to the reports directory. - - Returns: - str: Absolute path to model_manager/reports directory. - """ - # Get the directory where this file is located (model_manager/utils/repository/) - current_file_dir = path.dirname(path.abspath(__file__)) - # Navigate up to model_manager/ and then to reports/ - model_manager_dir = path.dirname(path.dirname(current_file_dir)) - reports_dir = path.join(model_manager_dir, 'reports') - return reports_dir diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index e381ea8..6b38749 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -161,14 +161,15 @@ async def main(): logger.custom_info(f'SDK metrics server initialized on port {SDK_METRICS_PORT}', metadata) + namespace = os.getenv('TEMPORAL_NAMESPACE', 'model-manager') temporal_client = await client.Client.connect( target_host=host, - namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), + namespace=namespace, runtime=new_runtime, tls=use_tls, ) - logger.custom_info(f'Temporal client initialized at {host}', metadata) + logger.custom_info(f'Temporal client initialized at {host}/{namespace}', metadata) # Create cleanup schedule (idempotent - only creates if doesn't exist) try: diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index 4ebb3cb..4e6c676 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -93,6 +93,8 @@ class TrainModel: Raises: ValueError: If experiment_run_id is missing or invalid """ + workflow.logger.info(f'Starting train_model workflow for {input_data}') + experiment_run_id = self._validate_experiment_run_id(input_data) input_data = {**input_data, 'experiment_run_id': experiment_run_id} diff --git a/requirements-dev.txt b/requirements-dev.txt index 50bf389..823c668 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -17,3 +17,4 @@ pytest-asyncio>=0.21.0 # Async test support (already in main requirements) # Development Tools ipython>=8.12.0 # Enhanced Python shell ipdb>=0.13.13 # IPython debugger +ipykernel<=6.26.0 # IPython kernel diff --git a/scripts/run_training_test.py b/scripts/run_training_test.py index 2cb3bb8..c1c13d6 100644 --- a/scripts/run_training_test.py +++ b/scripts/run_training_test.py @@ -32,6 +32,9 @@ import psycopg2 from dotenv import load_dotenv from psycopg2.extras import Json from temporalio import client +from dotenv import load_dotenv + +load_dotenv() # %% # --- configuration (edit here or use `.env` at repo root) --- @@ -54,9 +57,14 @@ PG = { TEMPORAL_HOST = os.getenv('TEMPORAL_HOST') TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE') -TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE') +TRAIN_TASK_QUEUE = "train_model-single-queue" TEMPORAL_TLS = os.getenv('TEMPORAL_USE_TLS', 'false').lower() in ('1', 'true', 'yes') + +print(MINIO_MC_ALIAS, MINIO_BUCKET, OBJECT_NAME, LOCAL_CSV) +print(PG) +print(TEMPORAL_HOST, TEMPORAL_NAMESPACE, TRAIN_TASK_QUEUE, TEMPORAL_TLS) + # %% # --- 1) database: delete previous row (same id), then insert `experiment_run` --- # Primary key column is `id` (see `experiment_tracking` updates). `request_data` matches the SQL sample in `input-sample.md`. @@ -97,7 +105,7 @@ with psycopg2.connect(**PG) as conn: EXPERIMENT_RUN_ID, 'test-experiment-name', 'test-run-name', - 'test-username', + 'vitor.santos@aignosi.com.br', 'ORCHESTRATOR_WAITING_PROC', None, now, @@ -142,15 +150,16 @@ _workflow_input = { 'opt_params': {}, } -# %% - c = await client.Client.connect( target_host=TH, namespace=TN, tls=TEMPORAL_TLS, ) + +# %% + wid = f'train-model-test-{uuid.uuid4()}' -await c.execute_workflow( # type: ignore[call-overload] +result = await c.execute_workflow( # type: ignore[call-overload] 'train_model', _workflow_input, id=wid, @@ -160,3 +169,6 @@ await c.execute_workflow( # type: ignore[call-overload] task_timeout=timedelta(minutes=5), ) print(wid) +print(result) + +# %% diff --git a/tests/activities/test_training.py b/tests/activities/test_training.py index 55caf2f..85cad71 100644 --- a/tests/activities/test_training.py +++ b/tests/activities/test_training.py @@ -93,7 +93,7 @@ def test_train_model_download_fails_notifies(training): 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } ) - training.minio_repository.download_file_sync = MagicMock(side_effect=OSError('minio')) + training.minio_repository.download_file = MagicMock(side_effect=OSError('minio')) training.send_notification = MagicMock() with pytest.raises(OSError, match='minio'): training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()}) @@ -129,7 +129,7 @@ def test_train_model_success_serializes_result(mock_mlflow, training): val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) - training.minio_repository.download_file_sync = MagicMock(return_value=b'csv') + training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w: setattr(x, 'mse_val', 0.1) or x @@ -185,7 +185,7 @@ def test_train_model_train_params_as_dict(mock_mlflow, training): tp = TrainModelParams.from_dict(d) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) - training.minio_repository.download_file_sync = MagicMock(return_value=b'csv') + training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w: x @@ -240,7 +240,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training): return b'val' raise AssertionError(object_name) - training.minio_repository.download_file_sync = MagicMock(side_effect=_dl) + training.minio_repository.download_file = MagicMock(side_effect=_dl) training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w: x @@ -272,7 +272,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training): training.mlflow_repository.start_run = _run_ctx training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) - assert training.minio_repository.download_file_sync.call_count == 2 + assert training.minio_repository.download_file.call_count == 2 mock_mlflow.log_artifact.assert_called() @@ -288,7 +288,7 @@ def test_train_model_value_error_when_paths_missing_after_report(training): val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) - training.minio_repository.download_file_sync = MagicMock(return_value=b'x') + training.minio_repository.download_file = MagicMock(return_value=b'x') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w: x diff --git a/tests/schedules/test_cleanup_schedule.py b/tests/schedules/test_cleanup_schedule.py index 63e9ef1..86f4d6a 100644 --- a/tests/schedules/test_cleanup_schedule.py +++ b/tests/schedules/test_cleanup_schedule.py @@ -125,6 +125,29 @@ async def test_schedule_exists_handles_exception(mock_temporal_client, mock_logg assert 'Error checking if schedule exists' in mock_logger.custom_error.call_args[0][0] +@pytest.mark.asyncio +async def test_needs_schedule_reconcile_handles_describe_exception(mock_logger, metadata): + """Test _needs_schedule_reconcile returns True and logs when describe fails.""" + from model_manager.schedules.cleanup_schedule import _needs_schedule_reconcile + + handle = AsyncMock() + handle.describe = AsyncMock(side_effect=RuntimeError('describe failed')) + + needs_reconcile = await _needs_schedule_reconcile( + schedule_handle=handle, + cleanup_task_queue='cleanup_files-model-manager-worker-queue', + logger=mock_logger, + metadata=metadata, + ) + + assert needs_reconcile is True + mock_logger.custom_error.assert_called_once() + assert ( + 'Error describing cleanup schedule for reconcile' + in mock_logger.custom_error.call_args[0][0] + ) + + # --- create_cleanup_schedule Tests --- diff --git a/tests/sientia/test_exceptions.py b/tests/sientia/test_exceptions.py new file mode 100644 index 0000000..c1d25e9 --- /dev/null +++ b/tests/sientia/test_exceptions.py @@ -0,0 +1,9 @@ +"""Unit tests for custom exception aliases.""" + +from mlflow.exceptions import MlflowException + +from model_manager.sientia.exceptions import SientiaMlException + + +def test_sientia_ml_exception_is_mlflow_exception_alias(): + assert SientiaMlException is MlflowException diff --git a/tests/utils/repository/test_data_manager_repository.py b/tests/utils/repository/test_data_manager_repository.py index 2573651..96bdf47 100644 --- a/tests/utils/repository/test_data_manager_repository.py +++ b/tests/utils/repository/test_data_manager_repository.py @@ -27,9 +27,11 @@ def test_train_test_split_dataframe_no_shuffle(): assert list(tr['a']) == [0, 1, 2, 3, 4] -def test_train_test_split_ndarray(): - arr = np.arange(20).reshape(10, 2) - tr, te = dmr.train_test_split(arr, train_size=0.5, shuffle=False, random_state=None) +def test_train_test_split_dataframe_returns_dataframes(): + df = pd.DataFrame(np.arange(20).reshape(10, 2), columns=['a', 'b']) + tr, te = dmr.train_test_split(df, train_size=0.5, shuffle=False, random_state=None) + assert isinstance(tr, pd.DataFrame) + assert isinstance(te, pd.DataFrame) assert tr.shape[0] == 5 and te.shape[0] == 5 diff --git a/tests/worker/test_prepare_worker.py b/tests/worker/test_prepare_worker.py index 0e0616d..c275c43 100644 --- a/tests/worker/test_prepare_worker.py +++ b/tests/worker/test_prepare_worker.py @@ -3,6 +3,12 @@ from unittest.mock import MagicMock, patch +def test_build_queue_name_without_runtime_uses_default_suffix(): + from model_manager.worker.prepare_worker import build_queue_name + + assert build_queue_name('TrainModel') == 'train_model-queue' + + def test_prepare_worker_train_queue_uses_train_limits(): from model_manager.worker.prepare_worker import prepare_worker from model_manager.workflows.train_model import TrainModel diff --git a/tests/workflows/test_train_model.py b/tests/workflows/test_train_model.py index 68bfd58..d3b4f94 100644 --- a/tests/workflows/test_train_model.py +++ b/tests/workflows/test_train_model.py @@ -290,9 +290,11 @@ async def test_run_training_failure_skips_cleanup_activity( @pytest.mark.asyncio -async def test_run_missing_experiment_run_id(): +@patch('model_manager.workflows.train_model.workflow') +async def test_run_missing_experiment_run_id(mock_wf): from model_manager.workflows.train_model import TrainModel + mock_wf.logger = Mock() with pytest.raises(ValueError, match='experiment_run_id is required'): await TrainModel().run({})