Issue #145: Created export_model_version.py

mlflow · Oct 1, 2023 · 05b3c2e · 05b3c2e
1 parent 050290c
commit 05b3c2e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/mlflow_export_import/model_version/export_model_version.py b/mlflow_export_import/model_version/export_model_version.py
@@ -52,11 +52,15 @@ def export_model_version(
 
     info_attr = {}
     mlflow_attr = { "model_version": _adjust_version(dict(vr)) }
+    msg = { "name": vr.name, "version": vr.version, "current_stage": vr.current_stage, "run_id": vr.run_id }
+    _logger.info(f"Exporting model verson: {msg}")
     io_utils.write_export_file(output_dir, "model_version.json", __file__, mlflow_attr, info_attr)
 
 
 def _adjust_version(vr):
-    """ Add nicely formatted timestamps and for aesthetic reasons order the dict attributes"""
+    """
+    Add nicely formatted timestamps and for aesthetic reasons order the dict attributes
+    """
     def _adjust_timestamp(dct, attr):
         dct[f"_{attr}"] = fmt_ts_millis(dct.get(attr))
     _adjust_timestamp(vr, "creation_timestamp")
@@ -70,7 +74,7 @@ def _adjust_timestamp(dct, attr):
 @opt_output_dir
 @opt_notebook_formats
 
-def main(model, 
+def main(model,
         version,
         output_dir,
         notebook_formats,
@@ -82,7 +86,7 @@ def main(model,
         model_name = model,
         version = version,
         output_dir = output_dir,
-        notebook_formats = notebook_formats,
+        notebook_formats = notebook_formats
     )
 
 

diff --git a/mlflow_export_import/run/export_run.py b/mlflow_export_import/run/export_run.py
@@ -52,6 +52,7 @@ def export_run(
     try:
         run = mlflow_client.get_run(run_id)
         if run.info.lifecycle_stage == "deleted" and not export_deleted_runs:
+            _logger.warning(f"Not exporting run '{run.info.run_id} because its lifecycle_stage is '{run.info.lifecycle_stage}'")
             return False
         experiment_id = run.info.experiment_id
         msg = { "run_id": run.info.run_id, "lifecycle_stage": run.info.lifecycle_stage, "experiment_id": run.info.experiment_id }
@@ -80,7 +81,7 @@ def export_run(
                run_id = run.info.run_id,
                dst_path = _filesystem.mk_local_path(dst_path),
                tracking_uri = mlflow_client._tracking_client.tracking_uri)
-        notebook = tags.get(MLFLOW_DATABRICKS_NOTEBOOK_PATH, None)
+        notebook = tags.get(MLFLOW_DATABRICKS_NOTEBOOK_PATH)
 
         # export notebook as artifact
         if notebook is not None:
@@ -115,7 +116,7 @@ def _get_metrics_with_steps(mlflow_client, run):
 def _export_notebook(dbx_client, output_dir, notebook, notebook_formats, run, fs):
     notebook_dir = os.path.join(output_dir, "artifacts", "notebooks")
     fs.mkdirs(notebook_dir)
-    revision_id = run.data.tags.get(MLFLOW_DATABRICKS_NOTEBOOK_REVISION_ID, None)
+    revision_id = run.data.tags.get(MLFLOW_DATABRICKS_NOTEBOOK_REVISION_ID)
     if not revision_id:
         _logger.warning(f"Cannot download notebook '{notebook}' for run '{run.info.run_id}' since tag '{MLFLOW_DATABRICKS_NOTEBOOK_REVISION_ID}' does not exist. Notebook is probably a Git Repo notebook.")
         return