gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1094

Browse Source 
Refactor activity imports and remove unused base and logger files; update requirements for library versioning
This commit is contained in:
vitor-aignosi
2025-06-09 11:08:07 -03:00
parent 11f41f358a
commit 5326051714
19 changed files with 23 additions and 908 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
laborious/activities/activities.py
View File

@@ -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):

26
laborious/activities/base.py
View File

@@ -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']

2
laborious/activities/gates.py
View File

@@ -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 (

4
laborious/activities/mlflow.py
View File

@@ -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):

2
laborious/activities/opc.py
View File

@@ -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

181
laborious/activities/postgres.py
View File

@@ -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()

22
laborious/utils/logger.py
View File

@@ -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

9
laborious/utils/policies.py
View File

@@ -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
)

241
laborious/utils/repository/model_repository.py
View File

@@ -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 {

5
laborious/worker/worker.py
View File

@@ -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():

2
laborious/workflows/predictions_batch.py
View File

@@ -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

2
laborious/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -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")

2
laborious/workflows/sub_workflows/prediction_process.py
View File

@@ -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

2
requirements.txt
View File

@@ -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

6
tests/laborious/activities/test_activities.py
View File

@@ -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']

35
tests/laborious/activities/test_base.py
View File

@@ -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"

159
tests/laborious/activities/test_postgres.py
View File

@@ -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()

190
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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'

37
tests/laborious/utils/test_logger.py
View File

@@ -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