refactor(prod_model): wrapping, removal of unnecessary code

Author: PLtier

github_dagger_workflow_project/04_prod_model.py

@@ -8,36 +8,40 @@
 model_name = "lead_model"
 experiment_best = pd.read_pickle(PROD_BEST_EXPERIMENT_PATH)
 
-prod_model = [
-    model
-    for model in client.search_model_versions(f"name='{model_name}'")
-    if dict(model)["current_stage"] == "Production"
-]
-prod_model_exists = len(prod_model) > 0
 
-if prod_model_exists:
-    prod_model_version = dict(prod_model[0])["version"]
-    prod_model_run_id = dict(prod_model[0])["run_id"]
+def get_production_model(model_name):
+    return [
+        model
+        for model in client.search_model_versions(f"name='{model_name}'")
+        if dict(model)["current_stage"] == "Production"
+    ]
+
+
+def get_model_score(run_id):
+    data, _ = mlflow.get_run(run_id)
+    return data[1]["metrics.f1_score"]
+
+
+def register_and_wait_model(run_id, artifact_path, model_name):
+    model_uri = f"runs:/{run_id}/{artifact_path}"
+    model_details = mlflow.register_model(model_uri=model_uri, name=model_name)
+    utils.wait_until_ready(model_details.name, model_details.version)
+    return dict(model_details)
+
 
 train_model_score = experiment_best["metrics.f1_score"]
-model_details = {}
-model_status = {}
 run_id = None
+prod_model = get_production_model(model_name)
+prod_model_exists = len(prod_model) > 0
 
 if prod_model_exists:
-    data, details = mlflow.get_run(prod_model_run_id)
-    prod_model_score = data[1]["metrics.f1_score"]
-
-    model_status["current"] = train_model_score
-    model_status["prod"] = prod_model_score
+    prod_model_run_id = dict(prod_model[0])["run_id"]
+    prod_model_score = get_model_score(prod_model_run_id)
 
     if train_model_score > prod_model_score:
         run_id = experiment_best["run_id"]
 else:
     run_id = experiment_best["run_id"]
 
 if run_id is not None:
-    model_uri = "runs:/{run_id}/{artifact_path}".format(run_id=run_id, artifact_path=artifact_path)
-    model_details = mlflow.register_model(model_uri=model_uri, name=model_name)
-    utils.wait_until_ready(model_details.name, model_details.version)
-    model_details = dict(model_details)
+    model_details = register_and_wait_model(run_id, artifact_path, model_name)