From 5326051714fc6cc0671720cfb8c2e2fc34b72190 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 9 Jun 2025 11:08:07 -0300 Subject: [PATCH] SIENTIAPDE-1094 Refactor activity imports and remove unused base and logger files; update requirements for library versioning --- laborious/activities/activities.py | 4 +- laborious/activities/base.py | 26 -- laborious/activities/gates.py | 2 +- laborious/activities/mlflow.py | 4 +- laborious/activities/opc.py | 2 +- laborious/activities/postgres.py | 181 ------------- laborious/utils/logger.py | 22 -- laborious/utils/policies.py | 9 - .../utils/repository/model_repository.py | 241 +----------------- laborious/worker/worker.py | 5 +- laborious/workflows/predictions_batch.py | 2 +- .../format_and_export_prediction.py | 2 +- .../sub_workflows/prediction_process.py | 2 +- requirements.txt | 2 +- tests/laborious/activities/test_activities.py | 6 +- tests/laborious/activities/test_base.py | 35 --- tests/laborious/activities/test_postgres.py | 159 ------------ .../utils/repository/test_model_repository.py | 190 +------------- tests/laborious/utils/test_logger.py | 37 --- 19 files changed, 23 insertions(+), 908 deletions(-) delete mode 100644 laborious/activities/base.py delete mode 100644 laborious/activities/postgres.py delete mode 100644 laborious/utils/logger.py delete mode 100644 laborious/utils/policies.py delete mode 100644 tests/laborious/activities/test_base.py delete mode 100644 tests/laborious/activities/test_postgres.py delete mode 100644 tests/laborious/utils/test_logger.py diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index e42e847..d3c9436 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -1,13 +1,13 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.postgres import Postgres + from sientia_do.temporal.activities.postgres import Postgres + from sientia_do.notifications.handlers import NotificationHandler from laborious.activities.mlflow import MLFlow from laborious.activities.gates import Gates from laborious.activities.opc import OPC from typing import Any from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler class Activities(Postgres, MLFlow, Gates, OPC): diff --git a/laborious/activities/base.py b/laborious/activities/base.py deleted file mode 100644 index 3adb0e4..0000000 --- a/laborious/activities/base.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any -from logging import Logger -from temporalio import activity -from sientia_do.notifications.handlers import NotificationHandler - - -class BaseActivity: - def __init__(self, logger: Logger, notification_handler: NotificationHandler): - self.logger = logger - self.notification_handler = notification_handler - - @activity.defn(name="prepare_activity") - async def prepare_activity(self, input_data: dict[str, Any]): - """ - Prepare the activity for the notification handler. - - Args: - workflow_name (str): The name of the workflow. - schedule_name (str): The name of the schedule. - model_name (str): The name of the model. - model_id (str): The id of the model. - """ - self.notification_handler.base_notification.pipeline_name = input_data['workflow_name'] - self.notification_handler.base_notification.schedule_name = input_data['schedule_name'] - self.notification_handler.base_notification.model_name = input_data['model_name'] - self.notification_handler.base_notification.model_id = input_data['model_id'] diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 1cf3fb9..799fed5 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -6,7 +6,7 @@ with workflow.unsafe.imports_passed_through(): from logging import Logger from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.models import NotificationLevel - from laborious.activities.base import BaseActivity + from sientia_do.temporal.activities.base import BaseActivity from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter from typing import Any from laborious.utils.filters.conditional_filters import ( diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index ae141a2..a209ed1 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -4,11 +4,11 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.base import BaseActivity + from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.notifications.handlers import NotificationHandler from laborious.utils.repository.model_repository import MLFlowRepository from typing import Any from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler class MLFlow(BaseActivity): diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 0f02c43..997b7c2 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -5,7 +5,7 @@ with workflow.unsafe.imports_passed_through(): from logging import Logger from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.models import NotificationLevel - from laborious.activities.base import BaseActivity + from sientia_do.temporal.activities.base import BaseActivity from laborious.utils.repository.opc_repository import OpcRepository from typing import Any import traceback diff --git a/laborious/activities/postgres.py b/laborious/activities/postgres.py deleted file mode 100644 index 9fc0ff7..0000000 --- a/laborious/activities/postgres.py +++ /dev/null @@ -1,181 +0,0 @@ -import traceback -from temporalio import workflow, activity - -from laborious.activities.base import BaseActivity -with workflow.unsafe.imports_passed_through(): - from sqlalchemy import create_engine - from sqlalchemy.orm import sessionmaker - from sqlalchemy.pool import QueuePool - from psycopg2.pool import ThreadedConnectionPool - from pandas import read_sql_query, DataFrame - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - from sientia_do.notifications.models import NotificationLevel - from typing import Any - - -class Postgres(BaseActivity): - def __init__(self, host: str, port: int, - user: str, password: str, dbname: str, - min_connections: int, max_connections: int, - logger: Logger, notification_handler: NotificationHandler): - self.host = host - self.port = port - self.user = user - self.password = password - self.dbname = dbname - - # Create SQLAlchemy engine with connection pooling - self.engine = create_engine( - f'postgresql://{user}:{password}@{host}:{port}/{dbname}', - poolclass=QueuePool, - pool_size=min_connections, - max_overflow=max_connections - min_connections, - pool_pre_ping=True - ) - self.session_factory = sessionmaker(bind=self.engine) - - BaseActivity.__init__(self, logger, notification_handler) - - def close(self): - self.engine.dispose() - - def __del__(self): - self.close() - - @activity.defn(name="load_custom_query") - async def load_custom_query(self, query: str) -> dict[str, Any]: - """ - Loads data from a custom query. - - Args: - query (str): The query to load data from. - - Returns: - dict[str, dict]: The data from the query. - """ - self.logger.info(f"Fetching data from query: {query}") - - data = None - with self.session_factory() as session: - try: - data = read_sql_query(query, self.engine) - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_LOADING_CUSTOM_QUERY", - message=f"Error fetching data from query: {e}", - block="load_custom_query", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - return {} - finally: - session.close() - - if data is None: - return {} - - # Converts any datetime datatype columns to string - for col in data.select_dtypes(include=['datetime64']).columns: - data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S') - - self.logger.info(f"Fetched {len(data)} rows") - self.logger.debug(f"Data: \n{data.to_string()}") - - return data.to_dict() - - @activity.defn(name="repeat_last_prediction") - async def repeat_last_prediction(self, query_items: dict[str, str]): - """ - Repeats the last prediction for a given model. - - Args: - query_items (dict[str, str]): The query items. Contains: - schema (str): The schema of the table. - table_name (str): The name of the table. - model (int): The model to repeat the prediction for. - - Returns: - None - """ - schema = query_items["schema"] - table_name = query_items["table_name"] - model = query_items["model"] - - repeat_query = f""" - INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at) - SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW() - FROM \"{schema}\".{table_name} - WHERE model_id = {model} - ORDER BY timestamp DESC - LIMIT 1; - """ - self.logger.info(f"Repeating last prediction for model {model}") - self.logger.debug(f"Query: {repeat_query}") - - with self.session_factory() as session: - try: - session.execute(repeat_query) - session.commit() - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_REPEATING_LAST_PREDICTION", - message=f"Error repeating last prediction: {e}", - block="repeat_last_prediction", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - finally: - session.close() - - @activity.defn(name="export_data_to_postgres") - async def export_data_to_postgres(self, input_data: dict[str, Any]): - """ - Exports data to a postgres table. - - Args: - input_data (dict[str, Any]): The data to export. Contains: - schema (str): The schema of the table. - table_name (str): The name of the table. - data (DataFrame): The data to export. - """ - - self.logger.debug( - f"Exporting data to postgres: {input_data['data']}") - - schema = input_data["schema"] - table_name = input_data["table_name"] - data = DataFrame(input_data["data"]) - - with self.session_factory() as session: - try: - data.to_sql(table_name, self.engine, schema=schema, - if_exists="append", index=False) - session.commit() - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES", - message=f"Error exporting data to postgres: {e}", - block="export_data_to_postgres", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - else: - self.logger.debug("Data exported to postgres") - finally: - session.close() diff --git a/laborious/utils/logger.py b/laborious/utils/logger.py deleted file mode 100644 index 42a9cfd..0000000 --- a/laborious/utils/logger.py +++ /dev/null @@ -1,22 +0,0 @@ -from os import getenv -import logging -import sys - - -def get_logger(name: str): - log_level = getenv('LOG_LEVEL', 'INFO').upper() - - logger = logging.getLogger(name) - logger.setLevel(log_level) - stream_handler = logging.StreamHandler(sys.stdout) - stream_handler.setLevel(log_level) - - stream_handler.setFormatter( - logging.Formatter( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - ) - - logger.addHandler(stream_handler) - - return logger diff --git a/laborious/utils/policies.py b/laborious/utils/policies.py deleted file mode 100644 index 8c7449a..0000000 --- a/laborious/utils/policies.py +++ /dev/null @@ -1,9 +0,0 @@ -from datetime import timedelta -from temporalio.common import RetryPolicy - -retry_policy = RetryPolicy( - initial_interval=timedelta(seconds=1), - backoff_coefficient=2.0, - maximum_interval=timedelta(minutes=1), - maximum_attempts=1 -) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index f2eaeb9..ed5936b 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1,16 +1,17 @@ """ Model Monitoring Repository -This module contains the ModelMonitoringRepository class, which is responsible for handling the communication with the Model Monitoring API. +This module contains the ModelMonitoringRepository class, which is responsible +for handling the communication with the Model Monitoring API. -It includes the methods that are used to answer ModelMonitoringService requests using the Model Monitoring API functions. +It includes the methods that are used to answer ModelMonitoringService requests using +the Model Monitoring API functions. By Monitoring we mean the evaluation of the performance of models, the generation of reports. """ from datetime import datetime import traceback -import mlflow import pandas as pd from sientia.ModelServing import ModelServing @@ -21,240 +22,6 @@ class MLFlowRepository(): self.model_serving = ModelServing(tracking_uri=host, username=username, password=password) - def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str): - """ - Get the current data as a DataFrame and update the prediction and target columns - - Parameters: - current_data (pd.DataFrame): the current data - model_name (str): the name of the model - target (str): the target column - - Returns: - DataFrame: the current data as a DataFrame - - - """ - predictions = current_data['prediction'] - - target = current_data[target] - current_data = self.model_serving.get_transformed_data( - model_name, current_data, by='model') - current_data['prediction'] = predictions - current_data['target'] = target - - return pd.DataFrame(current_data).dropna() - - def get_artifact(self, destination: str, search_by: str, run_id: str = None, - model_name: str = None, artifact_name: str = None) -> None: - """ - Get an artifact in MLflow by experiment or model and save it to a destination path using API. - If the artifact is searched by model, the latest production version will be used. - - Args: - destination: The destination path to save the artifact. - search_by: The way to search for the artifact ('experiment' or 'model'). - run_id: The run ID of the experiment (if search_by is "experiment"). - model_name: The name of the model (if search_by is "model"). - artifact_name: The path of the artifact to download. - - Returns: - artifact: The artifact(.csv) downloaded from MLflow. - """ - - self.model_serving.get_artifact(destination=destination, search_by=search_by, - run_id=run_id, model_name=model_name, artifact_name=artifact_name) - - def calculate_model_metrics(self, real_data, predictions, flag): - """ - Function to calculate the metrics of a model using API - - Parameters: - real_data (array): the real data - predictions (array): the predictions - - Returns: - dict: the metrics of the model including MSE and R2 - """ - return self.model_serving.get_model_metrics(reference_data=None, real_data=real_data, predictions=predictions, type_flag=flag) - - def get_experiment_by_run_id(self, run_id: str) -> dict: - # Get the run information using the run_id - run = mlflow.get_run(run_id) - - # Extract the experiment ID from the run - experiment_id = run.info.experiment_id - - # Get the experiment details using the experiment ID - experiment = mlflow.get_experiment(experiment_id) - experiment_name = experiment.name - return experiment_name - - def get_next_run_name(self, model_name: str) -> str: - """ - Function to get the next run number of a specific model - - Parameters: - model_name (str): the name of the model - - Returns: - str: the next run number - """ - - runs = mlflow.search_runs( - experiment_names=[model_name], order_by=["start_time desc"]) - next_run_number = len(runs) + 1 - return f"{model_name}-{next_run_number}" - - def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: - """ - Retrain a model with new data. - - Parameters: - data (pandas.DataFrame): The new data to use for retraining. - model_name (str): The name of the model to retrain. - metrics_list (list): The metrics to be used to compare the models. - compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one. - If False, the retrain will always be considered. - split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets. - If False, the data will be used as a unique block for retraining. - update_report (bool): If True, a report will be created with the data of the retrained model. - update_transformation (bool): If True, the model will be updated in the MLflow tracking server. - update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server. - shuffle_data (bool): If True, the data will be shuffled before splitting. - model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'. - - - Returns: - mlflow.sklearn.Model: The retrained prediction model. - mlflow.sklearn.Model: The retrained data model. - mse (float): The mean squared error of the retrained model. - r2 (float): The R-squared score of the retrained model. - """ - - # load predictor model - predictor_uri = f"models:/{model_name}/production" - # load transform model - latest_production_id = self.model_serving.get_model_run_id( - model_name, stage="Production" - ) - transform_uri = self.model_serving.get_model_uri( - latest_production_id, prediction=False - ) - # load - data_model = mlflow.sklearn.load_model(transform_uri) - prediction_model = mlflow.sklearn.load_model(predictor_uri) - data_model = data_model.fit(data) - treated_data = data_model.predict(data) - # align target column with treated_data - target_name = data_model.target_variable - y = data[target_name] - treated_data = pd.merge( - treated_data, y, left_index=True, right_index=True) - prediction_model = prediction_model.fit(treated_data) - # Example usage - experiment = self.get_experiment_by_run_id(latest_production_id) - pred_model_atributes = vars(prediction_model) # load class attributes - data_model_atributes = vars(data_model) # load class attributes - mlflow.set_experiment(experiment) - experiment_description = "Retrain model {model_name} with new data" - current_run_name = self.get_next_run_name(experiment) - with mlflow.start_run( - run_name=current_run_name, description=experiment_description - ) as _run: - # update transfomation model - # fixed parameters - for name_atribute, val_atribute in pred_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # update prediction model - for name_atribute, val_atribute in data_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # dynamic parameters, including model itself - mlflow.sklearn.log_model(data_model, "data_model") - file_path = f"laborious/data/raw_data_{model_name}.csv" - data.to_csv( - f"laborious/data/raw_data_{model_name}.csv", index=True) - # log the data raw - mlflow.log_artifact(file_path) - - # dynamic parameters, including model itself - mlflow.sklearn.log_model(prediction_model, "prediction_model") - mlflow.log_param("retrain", True) - - return "Model retrained successfully", experiment - - def get_experiment(self, experiment_name: str) -> int: - experiment = mlflow.get_experiment_by_name(experiment_name) - - if experiment is None: - raise ValueError(f'Experiment {experiment_name} not found') - - return int(experiment.experiment_id) - - def get_experiment_last_run(self, experiment_id: int) -> str: - runs = mlflow.search_runs( - experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" - ) - - # Filtrar apenas as runs onde params.retrain == True - filtered_runs = runs[runs["params.retrain"] == 'True'] - - # Converter a coluna 'end_time' para datetime - filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) - - # Ordenar o DataFrame de forma descendente pela coluna 'end_time' - filtered_runs = filtered_runs.sort_values( - by='end_time', ascending=False) - - # Pegar a última run_id do DataFrame filtrado e ordenado - latest_run_id = filtered_runs.iloc[0]['run_id'] - - return latest_run_id - - def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict: - # Registrar o modelo - # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro. - # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso. - mlflow.register_model( - f"runs:/{run_id}/prediction_model", model_name) - - # Colocar a versão do modelo em produção - # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production' - client = mlflow.tracking.MlflowClient() - - # Obter a versão mais recente registrada do modelo - model_versions = client.get_registered_model( - model_name).latest_versions - max_version = max(model_versions, key=lambda x: int(x.version)).version - - # Mover a versão mais recente do modelo para o estágio de 'Production' - client.transition_model_version_stage( - name=model_name, - version=max_version, - stage="Production", - archive_existing_versions=True - ) - - return { - 'model_name': model_name, - 'version': max_version, - 'mlflow_run_id': run_id - } - - def update_production_model(self, experiment: str, model_name: str) -> dict: - - experiment_id = self.get_experiment(experiment) - run_id = self.get_experiment_last_run(experiment_id) - metadata = self.update_production_model_by_run_id(run_id, model_name) - - metadata['mlflow_experiment_id'] = experiment_id - - return metadata - def transform(self, model_name: str, data: pd.DataFrame, model_retention: int): try: return { diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 5b902a9..f10716e 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -1,22 +1,23 @@ from temporalio import workflow, client from temporalio.worker import Worker -import sys + with workflow.unsafe.imports_passed_through(): import os + import sys import asyncio from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.sub_workflows.prediction_process import PredictionProcess from laborious.workflows.sub_workflows.format_and_export_prediction import \ FormatAndExportPrediction from laborious.activities.activities import Activities - from laborious.utils.logger import get_logger from laborious.utils.connectors_config import ( build_postgres_config, build_mlflow_config, build_opc_config ) from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.temporal.utils.logger import get_logger async def main(): diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 425f59f..9c4ab61 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -3,7 +3,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): from laborious.activities.activities import Activities from typing import Any - from laborious.utils.policies import retry_policy + from sientia_do.temporal.utils.policies import retry_policy from datetime import timedelta diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 3b1fad7..68a5958 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -4,7 +4,7 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.activities import Activities from typing import Any from datetime import timedelta - from laborious.utils.policies import retry_policy + from sientia_do.temporal.utils.policies import retry_policy @workflow.defn(name="format_and_export_prediction") diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 3164639..df740b3 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -3,7 +3,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): from laborious.activities.activities import Activities from typing import Any - from laborious.utils.policies import retry_policy + from sientia_do.temporal.utils.policies import retry_policy from datetime import timedelta diff --git a/requirements.txt b/requirements.txt index 90e049b..5604fc7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,5 +3,5 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 3b2ef49..31f83e8 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -1,7 +1,7 @@ from pytest import mark from unittest.mock import patch, MagicMock, ANY +from sientia_do.temporal.activities.postgres import Postgres from laborious.activities.activities import Activities -from laborious.activities.postgres import Postgres from laborious.activities.mlflow import MLFlow from laborious.activities.gates import Gates from laborious.activities.opc import OPC @@ -139,9 +139,9 @@ async def test_prepare_activity(_mock_opc_init, await activities.prepare_activity(input_data) - assert activities.notification_handler.base_notification.pipeline_name == input_data[ + assert activities.notification_handler.base_notification.pipeline == input_data[ 'workflow_name'] - assert activities.notification_handler.base_notification.schedule_name == input_data[ + assert activities.notification_handler.base_notification.trigger == input_data[ 'schedule_name'] assert activities.notification_handler.base_notification.model_name == input_data[ 'model_name'] diff --git a/tests/laborious/activities/test_base.py b/tests/laborious/activities/test_base.py deleted file mode 100644 index 6978acb..0000000 --- a/tests/laborious/activities/test_base.py +++ /dev/null @@ -1,35 +0,0 @@ -from unittest.mock import MagicMock -from laborious.activities.base import BaseActivity -from pytest import fixture, mark -from sientia_do.notifications.models import Notification - - -@fixture -def base_activity(): - return BaseActivity( - logger=MagicMock(), - notification_handler=MagicMock(), - ) - - -@mark.asyncio -async def test_prepare_activity(base_activity): - base_activity.notification_handler.base_notification = Notification( - project="project", - pipeline="pipeline", - trigger="-", - model_name="-", - model_id="-", - ) - - await base_activity.prepare_activity({ - 'workflow_name': 'test_workflow', - 'schedule_name': 'test_schedule', - 'model_name': 'test_model', - 'model_id': 'test_model_id' - }) - - assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule" - assert base_activity.notification_handler.base_notification.model_name == "test_model" - assert base_activity.notification_handler.base_notification.model_id == "test_model_id" - assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow" diff --git a/tests/laborious/activities/test_postgres.py b/tests/laborious/activities/test_postgres.py deleted file mode 100644 index e4a4545..0000000 --- a/tests/laborious/activities/test_postgres.py +++ /dev/null @@ -1,159 +0,0 @@ -from unittest.mock import MagicMock, patch -from pytest import fixture, mark -import pandas as pd -from laborious.activities.postgres import Postgres - - -@fixture -@patch("laborious.activities.postgres.create_engine") -def postgres_activity(_mock_create_engine): - return Postgres( - host="localhost", - port=5432, - user="test_user", - password="test_password", - dbname="test_db", - min_connections=1, - max_connections=5, - logger=MagicMock(), - notification_handler=MagicMock() - ) - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_read_sql_query.return_value = None - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 0 - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]}) - mock_data['date'] = pd.to_datetime('2022-01-01') - - mock_read_sql_query.return_value = mock_data - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 3 - assert "column1" in result - assert "column2" in result - assert "date" in result - assert result['date'] == {0: '2022-01-01 00:00:00'} - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_success(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]}) - - mock_read_sql_query.return_value = mock_data - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 2 - assert "column1" in result - assert "column2" in result - postgres_activity.logger.info.assert_called() - - -@mark.asyncio -async def test_load_custom_query_error(postgres_activity): - query = "SELECT * FROM non_existent_table" - error_msg = "Table not found" - - with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)): - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 0 - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_repeat_last_prediction_success(postgres_activity): - query_items = { - "schema": "public", - "table_name": "predictions", - "model": 1 - } - - with patch("sqlalchemy.orm.session.Session.execute") as mock_execute: - await postgres_activity.repeat_last_prediction(query_items) - - mock_execute.assert_called_once() - postgres_activity.logger.info.assert_called() - - -@mark.asyncio -async def test_repeat_last_prediction_error(postgres_activity): - query_items = { - "schema": "public", - "table_name": "predictions", - "model": 1 - } - error_msg = "Database error" - - with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)): - await postgres_activity.repeat_last_prediction(query_items) - - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_export_data_to_postgres_success(postgres_activity): - input_data = { - "schema": "public", - "table_name": "test_table", - "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]}) - } - - with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql: - await postgres_activity.export_data_to_postgres(input_data) - - mock_to_sql.assert_called_once() - postgres_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_export_data_to_postgres_error(postgres_activity): - input_data = { - "schema": "public", - "table_name": "test_table", - "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]}) - } - error_msg = "Export failed" - - with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)): - await postgres_activity.export_data_to_postgres(input_data) - - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_close(postgres_activity): - postgres_activity.close() - - postgres_activity.engine.dispose.assert_called_once() - - -@mark.asyncio -async def test_del(postgres_activity): - postgres_activity.close = MagicMock() - postgres_activity.__del__() - - postgres_activity.close.assert_called_once() diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index a675edb..6d10c1c 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,14 +1,14 @@ from unittest.mock import ANY, MagicMock, patch import numpy as np -from pandas import DataFrame import pytest from laborious.utils.repository.model_repository import MLFlowRepository @pytest.fixture def mlflow_repository(): - with patch('laborious.utils.repository.model_repository.ModelServing', autospec=True) as MockModelServing: - mock_instance = MockModelServing.return_value + with patch('laborious.utils.repository.model_repository.ModelServing', + autospec=True) as mock_model_serving: + mock_instance = mock_model_serving.return_value mock_instance.get_transformed_data = MagicMock() repo = MLFlowRepository( @@ -19,190 +19,6 @@ def mlflow_repository(): return repo -def test_get_current_data_df(mlflow_repository): - current_data = { - 'prediction': [1, 3], - 'target': [1, 1], - } - mlflow_repository.model_serving.get_transformed_data.return_value = { - 'var1': [1, 2], - 'var2': [2, np.nan], - } - expected = DataFrame({ - 'var1': [1], - 'var2': [2], - 'prediction': [1], - 'target': [1], - }) - output = mlflow_repository.get_current_data_df(current_data, - 'model', 'target') - - mlflow_repository.model_serving.get_transformed_data.assert_called_once_with( - 'model', current_data, by='model') - - diff = output.compare(expected) - assert diff.empty - - -def test_get_artifact(mlflow_repository): - mlflow_repository.get_artifact( - 'destination', 'search_by', 'run_id', 'model', 'artifact' - ) - mlflow_repository.model_serving.get_artifact.assert_called_once_with( - destination='destination', - search_by='search_by', - run_id='run_id', - model_name='model', - artifact_name='artifact' - ) - - -def test_calculate_model_metrics(mlflow_repository): - mlflow_repository.model_serving.get_model_metrics.return_value = 'data' - real_data = 'real_data' - predictions = 'predictions' - flag = 'flag' - output = mlflow_repository.calculate_model_metrics( - real_data, predictions, flag - ) - mlflow_repository.model_serving.get_model_metrics.assert_called_once_with( - reference_data=None, - real_data=real_data, - predictions=predictions, - type_flag=flag - ) - assert output == 'data' - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_by_run_id(mlflow, mlflow_repository): - mlflow.get_run.return_value = MagicMock( - info=MagicMock( - experiment_id='0', - ) - ) - mlflow.get_experiment.return_value = MagicMock() - mlflow.get_experiment.return_value.name = 'test' - - output = mlflow_repository.get_experiment_by_run_id('0') - assert output == 'test' - mlflow.get_run.assert_called_once_with('0') - mlflow.get_experiment.assert_called_once_with('0') - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_next_run_name(mlflow, mlflow_repository): - mlflow.search_runs.return_value = [1, 2, 3] - output = mlflow_repository.get_next_run_name('run') - assert output == 'run-4' - mlflow.search_runs.assert_called_once_with( - experiment_names=['run'], - order_by=['start_time desc'], - ) - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_success(mlflow, mlflow_repository): - mlflow.get_experiment_by_name.return_value = MagicMock( - experiment_id='0') - - output = mlflow_repository.get_experiment('test') - - assert output == 0 - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_error(mlflow, mlflow_repository): - mlflow.get_experiment_by_name.return_value = None - - try: - mlflow_repository.get_experiment('test') - except ValueError as e: - assert str(e) == 'Experiment test not found' - else: - assert False - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_last_run(mlflow, mlflow_repository): - mlflow.search_runs.return_value = DataFrame({ - 'params.retrain': ['True', 'False', 'True', 'False'], - 'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'], - 'run_id': ['0', '1', '2', '3'], - }) - - output = mlflow_repository.get_experiment_last_run(0) - - mlflow.search_runs.assert_called_once_with( - experiment_ids=[0], - filter_string="", - output_format="pandas", - ) - - assert output == '2' - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_update_production_model_by_run_id(mlflow, mlflow_repository): - client_mock = MagicMock() - mlflow.tracking.MlflowClient.return_value = client_mock - - client_mock.get_registered_model.return_value = MagicMock( - latest_versions=[ - MagicMock(version='1'), - MagicMock(version='2'), - MagicMock(version='3'), - ] - ) - output = mlflow_repository.update_production_model_by_run_id('0', 'test') - - mlflow.register_model.assert_called_once_with( - "runs:/0/prediction_model", - 'test', - ) - - mlflow.tracking.MlflowClient.assert_called_once() - client_mock.get_registered_model.assert_called_once_with('test') - client_mock.transition_model_version_stage.assert_called_once_with( - name='test', - version='3', - stage='Production', - archive_existing_versions=True, - ) - - assert output == { - 'model_name': 'test', - 'version': '3', - 'mlflow_run_id': '0', - } - - -def test_update_production_model(mlflow_repository): - connector = mlflow_repository - - with patch.object(connector, 'get_experiment', - return_value='0') as get_experiment: - with patch.object(connector, 'get_experiment_last_run', - return_value='2') as get_experiment_last_run: - with patch.object(connector, 'update_production_model_by_run_id', - return_value={'model_name': 'test', 'version': '3', - 'mlflow_run_id': '0'}) as update_production_model_by_run_id: - - output = connector.update_production_model('0', 'test') - - get_experiment.assert_called_once_with('0') - get_experiment_last_run.assert_called_once_with('0') - update_production_model_by_run_id.assert_called_once_with( - '2', 'test') - - assert output == { - 'model_name': 'test', - 'version': '3', - 'mlflow_run_id': '0', - 'mlflow_experiment_id': '0', - } - - def test_transform_success(mlflow_repository): data = 'data' model_name = 'model' diff --git a/tests/laborious/utils/test_logger.py b/tests/laborious/utils/test_logger.py deleted file mode 100644 index cb68cb4..0000000 --- a/tests/laborious/utils/test_logger.py +++ /dev/null @@ -1,37 +0,0 @@ -import os -from unittest.mock import patch -import logging -import pytest -from laborious.utils.logger import get_logger - - -@pytest.fixture -def mock_env_vars(): - with patch.dict(os.environ, {}, clear=True): - yield - - -@pytest.mark.usefixtures("mock_env_vars") -@patch('laborious.utils.logger.logging.Formatter') -@patch('laborious.utils.logger.logging.StreamHandler') -def test_get_logger_defaults(mock_stream_handler, mock_formatter): - """Test logger creation with default settings""" - # Mock the StreamHandler and Formatter - - logger = get_logger('test_logger') - - # Verify logger settings - assert logger.name == 'test_logger' - assert logger.level == logging.INFO - - # Verify handler configuration - mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO') - mock_stream_handler.return_value.setFormatter.assert_called_once() - - # Verify formatter configuration - mock_formatter.assert_called_once_with( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - - # Verify handler was added to logger - assert len(logger.handlers) == 1