refactor: в блокноте research, использовать mlflow nested runs

research/research.py

@@ -44,14 +44,8 @@ mlflow_experiment_id: Optional[str] = None
 # ID эксперимента MLFlow, имеет приоритет над `mlflow_experiment_name`.
 mlflow_experiment_name: Optional[str] = 'Current price predicion for used cars'
 # Имя эксперимента MLFlow (ниже приоритетом, чем `mlflow_experiment_id`).
-mlflow_baseline_run_name: str = 'Baseline model'
-# Имя ноговго прогона MLFlow для baseline модели.
-mlflow_feateng_run_name: str = 'Model with engineered features'
-# Имя ноговго прогона MLFlow для модели, использующей дополнительные признаки
-mlflow_feateng_filtered_run_name: str = 'Model with filtered engineered features'
-# Имя ноговго прогона MLFlow для модели, использующей дополнительные признаки и фильтрацию признаков
-mlflow_optimized_feateng_filtered_run_name: str = 'Optimized model with filtered engineered features'
-# Имя ноговго прогона MLFlow для модели с оптимизированными гиперпараметрами, использующей дополнительные признаки и фильтрацию признаков
+mlflow_root_run_name: str = 'Models'
+# Имя корневого прогона MLFlow (остальные прогоны будут созданы блокнотом внутри этого, как nested)
 
 # %%
 from collections.abc import Sequence
@@ -102,6 +96,7 @@ if mlflow_registry_uri is not None:
 # %%
 if mlflow_do_log:
     mlflow_experiment = mlflow.set_experiment(experiment_name=mlflow_experiment_name, experiment_id=mlflow_experiment_id)
+    mlflow_root_run_id = None  # изменяется позже
 
 # %%
 DATA_PATH = (
@@ -282,21 +277,34 @@ def score_predictions(target_test, target_test_predicted):
 
 
 # %%
-# использует глобальные переменные mlflow_do_log, mlflow_experiment
+# использует глобальные переменные mlflow_do_log, mlflow_experiment, mlflow_root_run_name
 def mlflow_log_model(
     model,
     model_params,
     metrics,
     *,
-    run_name,
+    nested_run_name,
     model_signature=None,
     input_example=None,
     #pip_requirements=None,
     comment_file_path=None,
 ):
+    global mlflow_root_run_id
     if not mlflow_do_log:
         return
-    with mlflow.start_run(experiment_id=mlflow_experiment.experiment_id, run_name=run_name):
+    experiment_id = mlflow_experiment.experiment_id
+    start_run_root_kwargs_extra = {}
+    if mlflow_root_run_id is not None:
+        start_run_root_kwargs_extra['run_id'] = mlflow_root_run_id
+    else:
+        start_run_root_kwargs_extra['run_name'] = mlflow_root_run_name
+    with mlflow.start_run(experiment_id=experiment_id, **start_run_root_kwargs_extra) as root_run:
+        if root_run.info.status not in ('RUNNING',):
+            raise RuntimeError('Cannot get the root run to run')
+        if mlflow_root_run_id is None:
+            mlflow_root_run_id = root_run.info.run_id
+        # важно одновременно использовать nested=True и parent_run_id=...:
+        with mlflow.start_run(experiment_id=experiment_id, run_name=nested_run_name, nested=True, parent_run_id=mlflow_root_run_id):
             _ = mlflow.sklearn.log_model(
                 model,
                 'model',
@@ -392,7 +400,7 @@ mlflow_log_model(
     pipeline,
     model_params=model_params,
     metrics={k: float(v) for k, v in metrics.items()},
-    run_name=mlflow_baseline_run_name,
+    nested_run_name='Baseline model',
     model_signature=mlflow_model_signature,
     input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
     #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
@@ -545,7 +553,7 @@ mlflow_log_model(
     pipeline,
     model_params=model_params,
     metrics={k: float(v) for k, v in metrics.items()},
-    run_name=mlflow_feateng_run_name,
+    nested_run_name='Model with engineered features',
     model_signature=mlflow_model_signature,
     input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
     #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
@@ -665,7 +673,7 @@ mlflow_log_model(
     pipeline,
     model_params=model_params,
     metrics={k: float(v) for k, v in metrics.items()},
-    run_name=mlflow_feateng_filtered_run_name,
+    nested_run_name='Model with filtered engineered features',
     model_signature=mlflow_model_signature,
     input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
     #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
@@ -785,7 +793,7 @@ mlflow_log_model(
     pipeline,
     model_params=model_params,
     metrics={k: float(v) for k, v in metrics.items()},
-    run_name=mlflow_optimized_feateng_filtered_run_name,
+    nested_run_name='Optimized model with filtered engineered features',
     model_signature=mlflow_model_signature,
     input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
     #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),