SIENTIAPDE-1214

SIENTIAPDE-1214: Refactor MLFlow and model repository methods to use model_config dictionary

- Updated MLFlow class methods to accept model_config instead of model_retention for improved flexibility.
- Modified model_repository methods to handle model_config, extracting necessary parameters for transformation and prediction.
- Adjusted predictions_batch and prediction_process workflows to utilize model_config for better configuration management.
- Commented out the previous sientia-mlops-library dependency in requirements.txt for clarity.
This commit is contained in:
vitor-aignosi
2025-09-12 14:23:35 -03:00
parent 4a4043c354
commit 6df04b72e5
35 changed files with 5333 additions and 18 deletions

View File

@@ -25,7 +25,7 @@ class MLFlowRepository():
username=username, password=password,
logger=logger)
def transform(self, model_name: str, data: pd.DataFrame, model_retention: int):
def transform(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
"""
Transform data using a model.
@@ -39,11 +39,19 @@ class MLFlowRepository():
"""
try:
model_retention = model_config.get('model_retention', 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')
return {
'success': True,
'content': self.model_serving.get_cached_transform(
model_name, data, model_retention).to_dict()
model_name, data, model_retention, flavor,
compressed, retention_target, transform_keyword
).to_dict()
}
except Exception as e:
@@ -55,7 +63,7 @@ class MLFlowRepository():
}
}
def predict(self, model_name: str, data: pd.DataFrame, model_retention: int):
def predict(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
"""
Predict data using a model.
@@ -68,11 +76,17 @@ class MLFlowRepository():
- dict: A dictionary containing the predicted data.
"""
try:
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('predict_flavor', 'pyfunc')
compressed = model_config.get('compressed', False)
retention_target = model_config.get('retention_target', 'model')
input_index = data.index
start_time = datetime.now()
data = self.model_serving.get_cached_predict(
model_name, data, model_retention)
model_name, data, model_retention, flavor,
compressed, retention_target
)
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])