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

97
model_manager/activities/mlflow.py
View File

@@ -1,20 +1,19 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from datetime import datetime
from pandas import Timestamp, to_datetime
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.activities.base import BaseActivity
import traceback
from typing import Any
import numpy as np
from pandas import DataFrame, to_datetime
from sientia_do.formatters import create_sample_dict
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.formatters import create_sample_dict
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.utils.repository.model_repository import MLFlowRepository
from typing import Any
import numpy as np
from pandas import DataFrame
import traceback
class MLFlow(BaseActivity):
@@ -36,8 +35,15 @@ class MLFlow(BaseActivity):
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize MLFlow activities with server configuration.
@@ -52,18 +58,17 @@ class MLFlow(BaseActivity):
Raises:
Exception: If MLFlowRepository initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger
f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
)
@activity.defn(name="request_transform")
@activity.defn(name='request_transform')
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Transform input data using MLFlow models.
@@ -99,7 +104,7 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug("Raw input data:", metadata)
self.debug('Raw input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
@@ -108,14 +113,12 @@ class MLFlow(BaseActivity):
)
# Pivot data for model input format
data = data.pivot(
index='timestamp', columns='variable',
values='value')
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
self.debug("Processed input data:", metadata)
self.debug('Processed input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Request transformation from MLFlow model
@@ -124,16 +127,20 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.debug(
f"Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data transformed successfully", metadata)
self.info('Data transformed successfully', metadata)
return response_data
@activity.defn(name="request_predict")
@activity.defn(name='request_predict')
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Execute predictions using MLFlow models.
@@ -169,14 +176,15 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug(f"Input data for: \n {data.head(5).to_string()}", metadata)
self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
# Convert numpy.nan to None for model compatibility
data.replace(np.nan, None, inplace=True)
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
# Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict(
@@ -184,13 +192,15 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data predicted successfully", metadata)
self.info('Data predicted successfully', metadata)
return response_data
@activity.defn(name="retrain_model")
@activity.defn(name='retrain_model')
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Retrain MLFlow models with updated training data.
@@ -234,8 +244,7 @@ class MLFlow(BaseActivity):
data.drop(columns=['model_id'], inplace=True, errors='ignore')
data.drop(columns=['created_at'], inplace=True, errors='ignore')
data = data.pivot(index='timestamp', columns='variable',
values='value')
data = data.pivot(index='timestamp', columns='variable', values='value')
data.sort_index(inplace=True)
data.reset_index(inplace=True)
@@ -244,15 +253,10 @@ class MLFlow(BaseActivity):
try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name
data=data, model_name=model_name
)
return {
'status': retrain_output,
'timestamp': timestamp,
'experiment': experiment
}
return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
@@ -261,12 +265,12 @@ class MLFlow(BaseActivity):
message=f'Error retraining model {model_name}: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name="update_production_model")
@activity.defn(name='update_production_model')
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Update production model with newly trained model version.
@@ -311,12 +315,12 @@ class MLFlow(BaseActivity):
status = input_data['status']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata)
f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try:
response = self.model_monitoring_repository.update_production_model(
experiment=experiment,
model_name=model_name
experiment=experiment, model_name=model_name
)
report = DataFrame([response])
@@ -325,8 +329,7 @@ class MLFlow(BaseActivity):
report['timestamp'] = timestamp
report['status'] = status
self.info(
f'Production model {model_name} updated successfully', metadata)
self.info(f'Production model {model_name} updated successfully', metadata)
return report.to_dict()
except Exception as e:
@@ -337,7 +340,7 @@ class MLFlow(BaseActivity):
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e