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,34 +277,47 @@ 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):
-        _ = mlflow.sklearn.log_model(
-            model,
-            'model',
-            signature=model_signature,
-            input_example=input_example,
-            #pip_requirements=pip_requirements,
-        )
-        if model_params is not None:
-            _ = mlflow.log_params(model_params)
-        if metrics is not None:
-            _ = mlflow.log_metrics(metrics)
-        if (comment_file_path is not None) and comment_file_path.exists():
-            mlflow.log_artifact(str(comment_file_path))
+    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',
+                signature=model_signature,
+                input_example=input_example,
+                #pip_requirements=pip_requirements,
+            )
+            if model_params is not None:
+                _ = mlflow.log_params(model_params)
+            if metrics is not None:
+                _ = mlflow.log_metrics(metrics)
+            if (comment_file_path is not None) and comment_file_path.exists():
+                mlflow.log_artifact(str(comment_file_path))
 
 
 # %% [markdown]
@@ -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),