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

SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows

Browse Source 
This commit includes several changes:

- Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability.
- Improved error handling and logging in gates.py and mlflow.py.
- Added input validation and filtering in gates.py to ensure data quality.
- Enhanced prediction formatting and storage policy management in gates.py.
- Updated metrics.py to use consistent naming conventions and labels.
- Refactored connectors_config.py to use type hints and improve code clarity.
- Updated conditional and MLFlow filters for better data quality checks.
- Improved model repository logic for retraining and updating models.
- Enhanced worker.py to include SDK metrics and improved error handling.
- Refactored workflows for better modularity and error handling.
- Updated tests to reflect the changes and improve test coverage.
This commit is contained in:
Bruno Domingues
2025-10-01 17:28:57 -03:00
parent b102f79087
commit dfc190c818
24 changed files with 1482 additions and 1399 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
model_manager/utils/repository/model_repository.py
View File

@@ -10,22 +10,23 @@ 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 pandas as pd
import mlflow
from datetime import datetime
from os import makedirs, path, remove
import mlflow
import pandas as pd
from sientia.ModelServing import ModelServing
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
class MLFlowRepository():
class MLFlowRepository:
def __init__(self, host, username, password, logger: Logger):
self.model_serving = ModelServing(tracking_uri=host,
username=username, password=password,
logger=logger)
self.model_serving = ModelServing(
tracking_uri=host, username=username, password=password, logger=logger
)
self.logger = logger
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
@@ -40,33 +41,32 @@ class MLFlowRepository():
# Get type of first element of index
index_type = type(index[0])
self.logger.custom_info(f"Index type: {index_type}", metadata)
self.logger.custom_info(f'Index type: {index_type}', metadata)
message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}'
# Check if all in index are of the same type
if not all(isinstance(i, index_type) for i in index):
raise ValueError(
f"{message}")
raise ValueError(f'{message}')
# Check type and converts to DATETIME_FORMAT_WITH_TZ
if index_type == str:
if index_type is str:
# Validate format of string and return error if not valid
try:
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
except ValueError:
raise ValueError(
f"{message}")
except ValueError as e:
raise ValueError(f'{message}') from e
elif index_type == datetime or index_type == pd.Timestamp:
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
else:
raise ValueError(
f"{message}")
raise ValueError(f'{message}')
return data
def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
def transform(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Transform data using a model.
@@ -81,41 +81,42 @@ class MLFlowRepository():
try:
self.logger.custom_debug(
f"Data received for model transformation: {data.to_csv()}", metadata)
f'Data received for model transformation: {data.to_csv()}', metadata
)
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('transform_flavor', 'sklearn')
compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model')
transform_keyword = model_config.get(
'transform_function_keyword', 'predict')
transform_keyword = model_config.get('transform_function_keyword', 'predict')
transformed_data = self.model_serving.get_cached_transform(
model_name, data, model_retention, flavor,
compressed, retention_target, transform_keyword
model_name,
data,
model_retention,
flavor,
compressed,
retention_target,
transform_keyword,
)
self.logger.custom_debug(
f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
f'Data received from model transformation: {transformed_data.to_csv()}', metadata
)
transformed_data = self.detect_and_parse_datetime_index(
transformed_data, metadata)
transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)
return {
'success': True,
'content': transformed_data.to_dict()
}
return {'success': True, 'content': transformed_data.to_dict()}
except Exception as e:
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
def predict(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Predict data using a model.
@@ -137,31 +138,26 @@ class MLFlowRepository():
start_time = datetime.now()
self.logger.custom_debug(
f"Data received for model prediction: {data.to_csv()}", metadata)
f'Data received for model prediction: {data.to_csv()}', metadata
)
data = self.model_serving.get_cached_predict(
model_name, data, model_retention, flavor,
compressed, retention_target
model_name, data, model_retention, flavor, compressed, retention_target
)
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])
self.logger.custom_debug(
f"Data received from model prediction: {data.to_csv()}", metadata)
f'Data received from model prediction: {data.to_csv()}', metadata
)
data.index = input_index
data['response_time'] = (end_time - start_time).total_seconds()
return {
'success': True,
'content': data.to_dict()
}
return {'success': True, 'content': data.to_dict()}
except Exception as e:
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def get_experiment_by_run_id(self, run_id: str) -> dict:
@@ -190,10 +186,9 @@ class MLFlowRepository():
Returns:
str: The next run name in format 'model_name-run_number'
"""
runs = mlflow.search_runs(
experiment_names=[model_name], order_by=["start_time desc"])
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}"
return f'{model_name}-{next_run_number}'
def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
"""
@@ -217,14 +212,10 @@ class MLFlowRepository():
- experiment: MLFlow experiment name
"""
# load predictor model
predictor_uri = f"models:/{model_name}/production"
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
)
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)
@@ -233,20 +224,16 @@ class MLFlowRepository():
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
return prediction_model, data_model, experiment
def perform_model_retrain(self,
prediction_model,
data_model,
experiment: str,
model_name: str,
data: pd.DataFrame):
def perform_model_retrain(
self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
):
"""
Execute the complete model retraining process in MLFlow.
@@ -271,7 +258,7 @@ class MLFlowRepository():
"""
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
experiment_description = f"Retrain model {model_name} with new data"
experiment_description = f'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
@@ -279,32 +266,32 @@ class MLFlowRepository():
# update transfomation model
# fixed parameters
for name_atribute, val_atribute in pred_model_atributes.items():
if name_atribute != "model":
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":
if name_atribute != 'model':
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(data_model, "data_model")
mlflow.sklearn.log_model(data_model, 'data_model')
makedirs("temp", exist_ok=True)
makedirs('temp', exist_ok=True)
file_path = f"temp/raw_data_{model_name}.csv"
file_path = f'temp/raw_data_{model_name}.csv'
data.to_csv(file_path, 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)
mlflow.sklearn.log_model(prediction_model, 'prediction_model')
mlflow.log_param('retrain', True)
# clear temp file
if path.exists(file_path):
remove(file_path)
return "Model retrained successfully", experiment
return 'Model retrained successfully', experiment
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
@@ -325,10 +312,10 @@ class MLFlowRepository():
- status_message (str): Retraining operation status
- experiment_name (str): MLFlow experiment identifier
"""
prediction_model, data_model, experiment = self.create_model_experiment(
model_name, data)
prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
retrain_result = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data)
prediction_model, data_model, experiment, model_name, data
)
return retrain_result
def get_experiment(self, experiment_name: str) -> int:
@@ -374,22 +361,21 @@ class MLFlowRepository():
"""
runs = mlflow.search_runs(
experiment_ids=[experiment_id],
filter_string="", # Sem filtro no MLflow ainda
output_format="pandas"
filter_string='', # Sem filtro no MLflow ainda
output_format='pandas',
)
if not isinstance(runs, pd.DataFrame):
raise ValueError('Runs is not a pandas DataFrame')
# Filtrar apenas as runs onde params.retrain == True
filtered_runs = runs[runs["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)
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']
@@ -423,16 +409,14 @@ class MLFlowRepository():
# 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)
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
model_versions = client.get_registered_model(model_name).latest_versions
if not isinstance(model_versions, list):
raise ValueError('Model versions is not a list')
@@ -441,17 +425,10 @@ class MLFlowRepository():
# 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
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
}
return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
def update_production_model(self, experiment: str, model_name: str) -> dict:
"""