Files
sientia-dataops-laborious_t…/laborious/activities/mlflow.py
vitor-aignosi 10081b71d2 SIENTIAPDE-1081
Enhance documentation across multiple modules with detailed parameter descriptions and usage examples
2025-06-09 12:05:26 -03:00

92 lines
3.3 KiB
Python

import numpy as np
from pandas import DataFrame
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
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
class MLFlow(BaseActivity):
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
BaseActivity.__init__(self, logger, notification_handler)
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
)
@activity.defn(name="request_transform")
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Access MLFlow model to get the transformed data.
Args:
- input_data (dict): The input data. Contains:
- data (dict[str, Any]): The data to transform.
- model_name (str): The name of the model.
- model_retention (int): The retention time of the model, in minutes.
Returns:
dict[str, Any]: The transformed data.
"""
self.logger.info('Transforming data...')
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
self.logger.debug("Raw input data:")
self.logger.debug(data)
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.logger.debug("Processed input data:")
self.logger.debug(data)
response_data = self.model_monitoring_repository.transform(
model_name, data, model_retention)
self.logger.debug("Response data:")
self.logger.debug(response_data)
return response_data
@activity.defn(name="request_predict")
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Access MLFlow model to get the predicted data.
Args:
- input_data (dict): The input data. Contains:
- data (dict[str, Any]): The data to predict.
- model_name (str): The name of the model.
- model_retention (int): The retention time of the model, in minutes.
Returns:
dict[str, Any]: The predicted data.
"""
self.logger.info('Predicting data...')
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
self.logger.debug(data)
data.replace(np.nan, None, inplace=True)
response_data = self.model_monitoring_repository.predict(
model_name, data, model_retention)
self.logger.debug(response_data)
return response_data