рефакторинг блокнота research в плане использования MLFlow

research/research.py

@@ -28,14 +28,24 @@ data_aug_pickle_path: Optional[str] = None
 data_aug_pickle_relpath: str = 'cars.aug.pickle'
 # Путь к файлу (pickle) для сохранения очищенного датасета относительно директории данных `data`. Игнорируется, если установлен data_aug_pickle_path.
 
-mlflow_tracking_server_uri: str = 'http://localhost:5000'
-# URL tracking сервера MLFlow
-mlflow_registry_uri: Optional[str] = None
-# URL сервера registry MLFlow (если не указан, используется `mlflow_tracking_server_uri`)
+model_comment_path: Optional[str] = None
+# Полный путь к текстовому файлу с произвольным комментарием для сохранения в MLFlow как артефакт вместе с моделью. Если не установлен, используется `research/<comment_relpath>`.
+model_comment_relpath: str = 'comment.txt'
+# Путь к текстовому файлу с произвольным комментарием для сохранения в MLFlow как артефакт вместе с моделью относительно директории `research`. Игнорируется, если установлен comment_path.
 
-mlflow_experiment_name: str = 'Current price predicion for used cars'
-mlflow_experiment_new: bool = False
+mlflow_tracking_server_uri: str = 'http://localhost:5000'
+# URL tracking-сервера MLFlow.
+mlflow_registry_uri: Optional[str] = None
+# URL сервера registry MLFlow (если не указан, используется `mlflow_tracking_server_uri`).
+
+mlflow_do_log: bool = False
+# Записывать ли прогон (run) в MLFlow; если True, при каждом исполнении блокнота создаётся новый прогон с именем `mlflow_run_name`.
+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_run_name: str = 'Baseline model'
+# Имя нового прогона MLFlow (используется для создания нового прогона, если `mlflow_do_log` установлен в True).
 
 # %%
 import os
@@ -56,11 +66,18 @@ import sklearn.preprocessing
 # %%
 BASE_PATH = pathlib.Path('..')
 
+# %%
+MODEL_INOUT_EXAMPLE_SIZE = 0x10
+
 # %%
 mlflow.set_tracking_uri(mlflow_tracking_server_uri)
 if mlflow_registry_uri is not None:
     mlflow.set_registry_uri(mlflow_registry_uri)
 
+# %%
+if mlflow_do_log:
+    mlflow_experiment = mlflow.set_experiment(experiment_name=mlflow_experiment_name, experiment_id=mlflow_experiment_id)
+
 # %%
 DATA_PATH = (
     pathlib.Path(os.path.dirname(data_aug_pickle_path))
@@ -74,9 +91,9 @@ DATA_PATH = (
 # %%
 with open(
     (
-    data_aug_pickle_path
-    if data_aug_pickle_path is not None
-    else (DATA_PATH / data_aug_pickle_relpath)
+        data_aug_pickle_path
+        if data_aug_pickle_path is not None
+        else (DATA_PATH / data_aug_pickle_relpath)
     ),
     'rb',
 ) as input_file:
@@ -169,6 +186,16 @@ tuple(map(len, (df_target_train, df_target_test)))
 # %% [markdown]
 # ## Создание пайплайнов обработки признаков и обучения модели
 
+# %%
+#MODEL_PIP_REQUIREMENTS_PATH = BASE_PATH / 'requirements' / 'requirements-isolated-research-model.txt'
+
+# %% [markdown]
+# Сигнатура модели для MLFlow:
+
+# %%
+mlflow_model_signature = mlflow.models.infer_signature(model_input=df_orig_features, model_output=df_target)
+mlflow_model_signature
+
 # %% [markdown]
 # Пайплайн предобработки признаков:
 
@@ -205,9 +232,11 @@ pipeline = sklearn.pipeline.Pipeline([
     ('preprocess', preprocess_transformer),
     ('regress', regressor),
 ])
+pipeline
 
 # %%
-pipeline
+model_params = pipeline.get_params()
+model_params
 
 # %% [markdown]
 # ## Baseline модель
@@ -232,43 +261,21 @@ metrics = {
 metrics
 
 # %%
-MODEL_PIP_REQUIREMENTS_PATH = BASE_PATH / 'requirements.txt'
-MODEL_COMMENTS_FILE_PATH = BASE_PATH / 'comment.txt'
-
-# %%
-MODEL_INOUT_EXAMPLE_SIZE = 0x10
-
-# %%
-model_inout_example = (df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE), df_target.head(MODEL_INOUT_EXAMPLE_SIZE))
-
-# %%
-mlflow_model_signature = mlflow.models.infer_signature(model_input=model_inout_example[0], model_output=model_inout_example[1])
-
-# %%
-mlflow_model_signature
-
-# %%
-model_params = pipeline.get_params()
-
-# %%
-model_params
-
-# %%
-if mlflow_experiment_new:
-    experiment = mlflow.get_experiment(mlflow.create_experiment(mlflow_experiment_name))
-else:
-    experiment = mlflow.set_experiment(experiment_name=mlflow_experiment_name)
-
-# %%
-with mlflow.start_run(experiment_id=experiment.experiment_id, run_name=mlflow_run_name):
-    _ = mlflow.sklearn.log_model(
-        pipeline,
-        'model',
-        signature=mlflow_model_signature,
-        input_example=model_inout_example[0],
-        pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
-    )
-    _ = mlflow.log_params(model_params)
-    _ = mlflow.log_metrics({k: float(v) for k, v in metrics.items()})
-    if MODEL_COMMENTS_FILE_PATH.exists():
-        mlflow.log_artifact(str(MODEL_COMMENTS_FILE_PATH))
+if mlflow_do_log:
+    with mlflow.start_run(experiment_id=mlflow_experiment.experiment_id, run_name=mlflow_run_name):
+        _ = mlflow.sklearn.log_model(
+            pipeline,
+            'model',
+            signature=mlflow_model_signature,
+            input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
+            #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
+        )
+        _ = mlflow.log_params(model_params)
+        _ = mlflow.log_metrics({k: float(v) for k, v in metrics.items()})
+        comment_file_path = (
+            model_comment_path
+            if model_comment_path is not None
+            else (BASE_PATH / 'research' / model_comment_relpath)
+        )
+        if comment_file_path.exists():
+            mlflow.log_artifact(str(comment_file_path))