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

feat: add experiment_name to TrainModelResult and update training logic

Browse Source 
- Introduced experiment_name parameter in TrainModelResult to enhance tracking of training experiments.
- Updated the Training class to utilize run_name and experiment_name for improved MLflow run management.
This commit is contained in:
vitor-aignosi
2026-04-16 09:43:39 -03:00
parent c75a5921f3
commit 1e05ebe147
3 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
model_manager/activities/training.py
View File

@@ -283,17 +283,17 @@ class Training(SientiaMonitoring):
self.info(f'Starting MLflow run for {train_params.model_type}', metadata)
with self.mlflow_repository.start_run(
model_name=train_params.model_name,
run_name=None,
experiment_name=f'{train_params.model_name}_experiment',
run_name=train_result.run_name,
experiment_name=train_result.experiment_name,
tags=None,
metadata=metadata,
) as run_info:
train_result.run_name = run_info.run_name
train_result.run_id = run_info.run_id
self._persist_training_artifacts(train_result, train_params, wrapper, metadata)
return {
'run_name': train_result.run_name,
'experiment_name': train_result.experiment_name,
'run_id': train_result.run_id,
'run_dir': train_result.run_dir,
}

1
model_manager/utils/models/train_model_result.py
View File

@@ -44,6 +44,7 @@ class TrainModelResult:
equation: dict | None = None
equation_path: str | None = None
run_name: str | None = None
experiment_name: str | None = None
run_id: str | None = None
report_path: str | None = None
train_data_path: str | None = None

6
model_manager/utils/repository/data_manager_repository.py
View File

@@ -20,6 +20,7 @@ from os import makedirs, path
from shutil import rmtree
from typing import Any
from mlflow.entities import experiment
import numpy as np
import pandas as pd
from sientia_do.observability.logger import Logger
@@ -200,9 +201,10 @@ class DataManagerRepository(SientiaMonitoring):
metadata,
)
run_name = f'train_model_{params.model_type}_{params.model_name}_{params.experiment_run_id}'
experiment_name = f'train_model_{params.model_type}_{params.model_name}_{params.experiment_run_id}'
run_name = f'{experiment_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
return TrainModelResult(params=params, train_data=train_data, val_data=val_data, run_name=run_name)
return TrainModelResult(params=params, train_data=train_data, val_data=val_data, run_name=run_name, experiment_name=experiment_name)
def _as_series(self, pred: pd.DataFrame | pd.Series) -> pd.Series:
if isinstance(pred, pd.Series):