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

SIENTIAPDE-1273

Browse Source 
Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.
This commit is contained in:
vitor-aignosi
2025-11-13 16:40:23 -03:00
parent 6ac0f38d59
commit b68674fe64
15 changed files with 1639 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -217,6 +217,24 @@ class MLFlowRepository(SientiaMonitoring):
run_info = mlflow.get_run(run_id)
return run_info.data.params
def check_artifact_exists(self, run_id: str,
artifact_path: str, metadata: dict[str, Any]) -> bool:
"""
Check if an artifact exists in the MLflow Model Registry.
Args:
run_id (str): Run identifier to inspect.
artifact_path (str): Path to the artifact to check.
Returns:
bool: True if the artifact exists, False otherwise.
"""
artifacts = self.client.list_artifacts(run_id)
self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata)
self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata)
return any(artifact.path == artifact_path for artifact in artifacts)
"""
Functions related to download and load models
"""
@@ -263,7 +281,7 @@ class MLFlowRepository(SientiaMonitoring):
async def load_artifact_dataframe(self, model_name: str, artifact_path: str,
metadata: dict[str, Any]) -> pd.DataFrame:
metadata: dict[str, Any]) -> pd.DataFrame | None:
"""
Load the dataframe content of an artifact from the MLflow Model Registry.
@@ -277,9 +295,13 @@ class MLFlowRepository(SientiaMonitoring):
pd.DataFrame: The dataframe content of the artifact.
"""
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
artifact_path = path.join("runs:/", run_id, artifact_path)
core_labels = self.get_core_labels(metadata, operation_type='load_text')
if not self.check_artifact_exists(run_id, artifact_path, metadata):
return None
artifact_path = path.join("runs:/", run_id, artifact_path)
start_time = time.time()
try:
content = mlflow.artifacts.load_text(artifact_path)
@@ -290,6 +312,8 @@ class MLFlowRepository(SientiaMonitoring):
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
self.debug(f'Content of {run_id}/{artifact_path}: \n{content}', metadata)
dataframe = pd.read_csv(StringIO(content))
self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
@@ -680,6 +704,41 @@ class MLFlowRepository(SientiaMonitoring):
Functions related to model retraining
"""
def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame,
target_name: str) -> pd.DataFrame:
"""
Get prediction data from prediction model.
"""
input_index = retrain_dataset.index
prediction_data = prediction_model.predict(retrain_dataset)
if isinstance(prediction_data, pd.DataFrame):
prediction_data.columns = pd.Index(['prediction'])
else:
prediction_data = pd.DataFrame(prediction_data, columns=['prediction'])
prediction_data.index = input_index
# Merge prediction data with retrain_dataset on index
prediction_data = pd.merge(
retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')
# Rename column "target_name" to "target"
prediction_data.rename(columns={target_name: 'target'}, inplace=True)
prediction_data['timestamp'] = prediction_data.index
prediction_data.reset_index(drop=True, inplace=True)
prediction_data.sort_values(
by='timestamp', ascending=True, inplace=True
)
return prediction_data
async def fit_models(
self,
model_name: str,
@@ -689,7 +748,7 @@ class MLFlowRepository(SientiaMonitoring):
transform_flavor: str = 'sklearn',
predict_flavor: str = 'sklearn',
target_name: str | None = None,
) -> dict[str, dict[str, Any]]:
) -> dict[str, Any]:
"""
Prepare models and data for a retraining run.
@@ -799,6 +858,10 @@ class MLFlowRepository(SientiaMonitoring):
prediction_model.fit(retrain_dataset)
# get prediction data
prediction_data = self.get_prediction_data(
prediction_model, retrain_dataset, target_name)
self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
retrain_data = {
@@ -807,6 +870,7 @@ class MLFlowRepository(SientiaMonitoring):
'artifact_path': prediction_artifact_path,
},
'data_model': {'model': data_model, 'artifact_path': data_artifact_path},
'prediction_data': prediction_data,
}
return retrain_data
@@ -885,6 +949,7 @@ class MLFlowRepository(SientiaMonitoring):
prediction_model = retrain_data['prediction_model']
data_model = retrain_data['data_model']
prediction_data = retrain_data['prediction_data']
model_temp_path = path.join(ARTIFACTS_PATH, model_name)
@@ -908,10 +973,12 @@ class MLFlowRepository(SientiaMonitoring):
self.debug(f'Attributes: {retrain_params}', metadata)
data_path = f'{model_temp_path}/retrain_data.csv'
prediction_data_path = f'{model_temp_path}/evaluation_data.csv'
makedirs(model_temp_path, exist_ok=True)
data.to_csv(data_path, index=True)
data.to_csv(data_path, index=False)
prediction_data.to_csv(prediction_data_path, index=False)
self.info(
f'Starting model upload for {experiment_name} with run name {current_run_name}',
@@ -945,6 +1012,7 @@ class MLFlowRepository(SientiaMonitoring):
# log the data raw
mlflow.log_artifact(data_path)
mlflow.log_artifact(prediction_data_path)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
@@ -1372,59 +1440,3 @@ class MLFlowRepository(SientiaMonitoring):
return metadata_result
async def get_model_metrics(self,
analyse_data: pd.DataFrame,
train_reference_data: pd.DataFrame,
test_reference_data: pd.DataFrame,
target_name: str,
drift_metrics: list[str],
performance_metrics: list[str],
metadata: dict) -> dict[str, Any]:
"""
Calculates drift and preformance metrics for a model.
Args:
analyse_data (pd.DataFrame): The data to analyse.
train_reference_data (pd.DataFrame): The train reference data.
test_reference_data (pd.DataFrame): The test reference data.
target_name (str): The target name.
drift_metrics (list[str]): The drift metrics.
performance_metrics (list[str]): The performance metrics.
metadata (dict): The metadata.
Returns:
dict: The metrics.
"""
columns = train_reference_data.drop(
columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore').columns
config = {
"target": target_name,
"prediction": "prediction",
"timestamp": "timestamp",
"columns": columns,
}
model_analysis = ModelAnalysis(config=config)
core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
start_time = time.time()
try:
model_analysis.detect_univariate_drift(
reference_df=train_reference_data,
analyse_df=analyse_data,
features=columns,
timestamp_column=config['timestamp'],
metrics=drift_metrics,
chunk_period="s"
)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
return {}