Issue #43: implemented export/import of 'new' Run.inputs

Author: amesar

mlflow_export_import/run/export_run.py

@@ -67,7 +67,8 @@ def export_run(
             "info": info,
             "params": run.data.params,
             "metrics": _get_metrics_with_steps(mlflow_client, run),
-            "tags": tags
+            "tags": tags,
+            "inputs": _inputs_to_dict(run.inputs)
         }
         io_utils.write_export_file(output_dir, "run.json", __file__, mlflow_attr)
         fs = _filesystem.get_filesystem(".")
@@ -130,6 +131,15 @@ def _export_notebook(dbx_client, output_dir, notebook, notebook_formats, run, fs
     download_notebook(notebook_dir, notebook, revision_id, notebook_formats, dbx_client)
 
 
+def _inputs_to_dict(inputs):
+    def to_dict(ds):
+        return {
+            "dataset": utils.strip_underscores(ds.dataset),
+            "tags": [ utils.strip_underscores(tag) for tag in ds.tags ]
+        }
+    return [ to_dict(x) for x in inputs.dataset_inputs ]
+
+
 @click.command()
 @opt_run_id
 @opt_output_dir

mlflow_export_import/run/import_run.py

@@ -22,7 +22,7 @@
 from mlflow_export_import.common.filesystem import mk_local_path
 from mlflow_export_import.common import filesystem as _filesystem
 from mlflow_export_import.common import MlflowExportImportException
-from mlflow_export_import.client.http_client import create_dbx_client
+from mlflow_export_import.client.http_client import create_dbx_client, create_http_client
 from . import run_data_importer
 from . import run_utils
 
@@ -61,6 +61,7 @@ def _mk_ex(src_run_dct, dst_run_id, exp_name):
     }
 
     mlflow_client = mlflow_client or mlflow.MlflowClient()
+    http_client = create_http_client(mlflow_client)
     dbx_client = create_dbx_client(mlflow_client)
 
     _logger.info(f"Importing run from '{input_dir}'")
@@ -69,7 +70,6 @@ def _mk_ex(src_run_dct, dst_run_id, exp_name):
     src_run_path = os.path.join(input_dir, "run.json")
     src_run_dct = io_utils.read_file_mlflow(src_run_path)
     in_databricks = "DATABRICKS_RUNTIME_VERSION" in os.environ
-    #_logger.debug(f"in_databricks: {in_databricks}")
 
     run = mlflow_client.create_run(exp.experiment_id)
     run_id = run.info.run_id
@@ -83,6 +83,8 @@ def _mk_ex(src_run_dct, dst_run_id, exp_name):
             use_src_user_id,
             in_databricks
         )
+        _import_inputs(http_client, src_run_dct, run_id)
+
         path = os.path.join(input_dir, "artifacts")
         if os.path.exists(_filesystem.mk_local_path(path)):
             mlflow_client.log_artifacts(run_id, mk_local_path(path))
@@ -141,6 +143,12 @@ def _upload_databricks_notebook(dbx_client, input_dir, src_run_dct, dst_notebook
         _logger.warning(f"Cannot save notebook '{dst_notebook_path}'. {e}")
 
 
+def _import_inputs(http_client, src_run_dct, run_id):
+    inputs = src_run_dct.get("inputs")
+    dct = { "run_id": run_id, "datasets": inputs }
+    http_client.post("runs/log-inputs", dct)
+
+
 @click.command()
 @opt_input_dir
 @opt_experiment_name