Issue mlflow#190: export-experiment: fix for exporting run_id from ot…

…her experiment
matbun · Oct 3, 2024 · 91f0da5 · 91f0da5
1 parent f00ea57
commit 91f0da5
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/mlflow_export_import/experiment/export_experiment.py b/mlflow_export_import/experiment/export_experiment.py
@@ -131,6 +131,16 @@ def _export_run(mlflow_client, run_or_run_id, experiment_id, output_dir,
             _logger.info(f"Not exporting run: {msg}")
             return
 
+    if run.info.experiment_id != experiment_id:
+        msg = {
+           "run_id": {run.info.run_id},
+            "run.experiment_id": {run.info.experiment_id},
+            "experiment_id": {experiment_id}
+        }
+        _logger.warning(f"Not exporting run since it doesn't belong to experiment: {msg}")
+        failed_run_ids.append(run.info.run_id)
+        return
+
     is_success = export_run(
         run_id = run.info.run_id,
         output_dir = os.path.join(output_dir, run.info.run_id),

diff --git a/tests/open_source/test_experiments.py b/tests/open_source/test_experiments.py
@@ -240,3 +240,38 @@ def test_exp_with_multiple_runs_nonexistent_run(mlflow_context):
     runs2 = client2.search_runs(exp2.experiment_id)
     assert len(runs2) == 1
     compare_runs(mlflow_context, run1_ok, runs2[0])
+
+def test_exp_with_run_from_other_experiment(mlflow_context):
+    client1, client2 = mlflow_context.client_src, mlflow_context.client_dst
+    init_output_dirs(mlflow_context.output_dir)
+    exp1a = create_experiment(client1)
+    exp1b = create_experiment(client1)
+
+    mlflow.set_experiment(exp1a.name)
+    run1a = _create_simple_run(client1)
+
+    mlflow.set_experiment(exp1b.name)
+    run1b = _create_simple_run(client1)
+
+    runs1 = client1.search_runs(exp1a.experiment_id)
+    assert len(runs1) == 1
+
+    run_ids = [ run1b.info.run_id, run1a.info.run_id ]
+
+    export_experiment(
+        mlflow_client = client1,
+        experiment_id_or_name = exp1a.name,
+        run_ids = run_ids,
+        output_dir = mlflow_context.output_dir
+    )
+
+    exp_name2 = mk_dst_experiment_name(exp1a.name)
+    exp2 = import_experiment(
+        mlflow_client = client2,
+        experiment_name = exp_name2,
+        input_dir = mlflow_context.output_dir
+    )
+    exp2 = client2.get_experiment_by_name(exp_name2)
+    runs2 = client2.search_runs(exp2.experiment_id)
+    assert len(runs2) == 1
+    compare_runs(mlflow_context, run1a, runs2[0])