diff --git a/mlflow_export_import/experiment/export_experiment.py b/mlflow_export_import/experiment/export_experiment.py index 1748b5ee..a4ff4526 100644 --- a/mlflow_export_import/experiment/export_experiment.py +++ b/mlflow_export_import/experiment/export_experiment.py @@ -45,7 +45,6 @@ def export_experiment( """ mlflow_client = mlflow_client or create_mlflow_client() - dbx_client = create_dbx_client(mlflow_client) run_start_time_str = run_start_time if run_start_time: @@ -91,6 +90,7 @@ def export_experiment( mlflow_attr = { "experiment": exp_dct , "runs": ok_run_ids } if export_permissions: + dbx_client = create_dbx_client(mlflow_client) mlflow_attr["permissions"] = ws_permissions_utils.get_experiment_permissions(dbx_client, exp.experiment_id) io_utils.write_export_file(output_dir, "experiment.json", __file__, mlflow_attr, info_attr) diff --git a/mlflow_export_import/run/export_run.py b/mlflow_export_import/run/export_run.py index c0549888..eab0d7e8 100644 --- a/mlflow_export_import/run/export_run.py +++ b/mlflow_export_import/run/export_run.py @@ -47,7 +47,6 @@ def export_run( """ mlflow_client = mlflow_client or create_mlflow_client() - dbx_client = create_dbx_client(mlflow_client) if notebook_formats is None: notebook_formats = [] @@ -96,6 +95,7 @@ def export_run( # export notebook as artifact if notebook is not None: if len(notebook_formats) > 0: + dbx_client = create_dbx_client(mlflow_client) _export_notebook(dbx_client, output_dir, notebook, notebook_formats, run, fs) elif len(notebook_formats) > 0: _logger.warning(f"No notebooks to export for run '{run_id}' since tag '{MLFLOW_DATABRICKS_NOTEBOOK_PATH}' is not set.")