Up

Signed-off-by: Peter Jung <[email protected]>

Author: kongzii

mlflow_export_import/common/dump_run.py

@@ -9,8 +9,6 @@
 INDENT = "  "
 MAX_LEVEL = 1
 TS_FORMAT = "%Y-%m-%d_%H:%M:%S"
-client = mlflow.tracking.MlflowClient()
-print("MLflow Tracking URI:", mlflow.get_tracking_uri())
 
 def dump_run(run, max_level=1, indent=""):
     dump_run_info(run.info,indent)
@@ -27,17 +25,19 @@ def dump_run(run, max_level=1, indent=""):
     num_bytes, num_artifacts = dump_artifacts(run.info.run_id, "", 0, max_level, indent+INDENT)
     print(f"{indent}Total: bytes: {num_bytes} artifacts: {num_artifacts}")
     return run, num_bytes, num_artifacts
-        
+
 def dump_run_id(run_id, max_level=1, indent=""):
+    client = mlflow.tracking.MlflowClient()
     run = client.get_run(run_id)
     return dump_run(run,max_level,indent)
 
 def dump_run_info(info, indent=""):
     print("{}RunInfo:".format(indent))
+    client = mlflow.tracking.MlflowClient()
     exp = client.get_experiment(info.experiment_id)
     if exp is None:
         print(f"ERROR: Cannot find experiment ID '{info.experiment_id}'")
-        return 
+        return
     print("{}  name: {}".format(indent,exp.name))
     for k,v in sorted(info.__dict__.items()):
         if not k.endswith("_time"):
@@ -58,8 +58,9 @@ def _dump_time(info, k, indent=""):
     return v
 
 def dump_artifacts(run_id, path, level, max_level, indent):
-    if level+1 > max_level: 
+    if level+1 > max_level:
         return 0,0
+    client = mlflow.tracking.MlflowClient()
     artifacts = client.list_artifacts(run_id,path)
     num_bytes, num_artifacts = (0,0)
     for j,art in enumerate(artifacts):

mlflow_export_import/common/find_artifacts.py

@@ -7,16 +7,15 @@
 import click
 import mlflow
 
-client = mlflow.tracking.MlflowClient()
-print("MLflow Tracking URI:", mlflow.get_tracking_uri())
 
 def find_artifacts(run_id, path, target, max_level=sys.maxsize):
     return _find_artifacts(run_id, path, target, max_level, 0, [])
 
 def _find_artifacts(run_id, path, target, max_level, level, matches):
-    if level+1 > max_level: 
+    if level+1 > max_level:
         return matches
-    artifacts = client.list_artifacts(run_id,path)
+    client = mlflow.tracking.MlflowClient()
+    artifacts = client.list_artifacts(run_id, path)
     for art in artifacts:
         #print(f"art_path: {art.path}")
         filename = os.path.basename(art.path)
@@ -40,5 +39,5 @@ def main(run_id, path, target, max_level): # pragma: no cover
     for x in matches:
         print(" ",x)
 
-if __name__ == "__main__": 
+if __name__ == "__main__":
     main()

mlflow_export_import/model/export_model.py

@@ -11,18 +11,19 @@
 from mlflow_export_import import utils, click_doc
 
 class ModelExporter():
-    def __init__(self,  mlflow_client, export_source_tags=False, notebook_formats=None, stages=None, versions=None, export_run=True):
+    def __init__(self,  mlflow_client, export_source_tags=False, notebook_formats=None, stages=None, versions=None, export_run=True, host=None):
         """
         :param mlflow_client: MLflow client or if None create default client.
         :param export_source_tags: Export source run metadata tags.
         :param notebook_formats: List of notebook formats to export. Values are SOURCE, HTML, JUPYTER or DBC.
         :param stages: Stages to export. Default is all stages. Values are Production, Staging, Archived and None.
         :param versions: Versions to export. Default is all versions. Values are valid integer numbers.
         :param export_run: Export the run that generated a registered model's version.
+        :param host: Pass host to the MlflowHttpClient and RunExporter.
         """
         self.mlflow_client = mlflow_client
-        self.http_client = MlflowHttpClient()
-        self.run_exporter = RunExporter(self.mlflow_client, export_source_tags=export_source_tags, notebook_formats=notebook_formats)
+        self.http_client = MlflowHttpClient(host=host)
+        self.run_exporter = RunExporter(self.mlflow_client, export_source_tags=export_source_tags, notebook_formats=notebook_formats, host=host)
         self.stages = self._normalize_stages(stages)
         self.versions = self._normalize_versions(versions)
         self.export_run = export_run
@@ -66,7 +67,7 @@ def _export_model(self, model_name, output_dir):
                 run = self.mlflow_client.get_run(run_id)
                 dct = dict(vr)
                 dct["_run_artifact_uri"] = run.info.artifact_uri
-                experiment = mlflow.get_experiment(run.info.experiment_id)
+                experiment = self.mlflow_client.get_experiment(run.info.experiment_id)
                 dct["_experiment_name"] = experiment.name
                 model["registered_model"]["latest_versions"].append(dct)
                 exported_versions += 1

mlflow_export_import/model/import_model.py

@@ -13,14 +13,15 @@
 
 class BaseModelImporter():
     """ Base class of ModelImporter subclasses. """
-    def __init__(self, mlflow_client, run_importer=None, await_creation_for=None):
+    def __init__(self, mlflow_client, run_importer=None, await_creation_for=None, host=None):
         """
         :param mlflow_client: MLflow client or if None create default client.
         :param run_importer: RunImporter instance.
         :param await_creation_for: Seconds to wait for model version crreation.
+        :param host: Pass host to the RunImporter.
         """
         self.mlflow_client = mlflow_client
-        self.run_importer = run_importer if run_importer else RunImporter(self.mlflow_client, mlmodel_fix=True)
+        self.run_importer = run_importer if run_importer else RunImporter(self.mlflow_client, mlmodel_fix=True, host=host)
         self.await_creation_for = await_creation_for
 
     def _import_version(self, model_name, src_vr, dst_run_id, dst_source, sleep_time):
@@ -77,27 +78,29 @@ def _import_model(self, model_name, input_dir, delete_model=False, verbose=False
 
 class ModelImporter(BaseModelImporter):
     """ Low-level 'point' model importer  """
-    def __init__(self, mlflow_client, run_importer=None, await_creation_for=None):
-        super().__init__(mlflow_client, run_importer, await_creation_for=await_creation_for)
+    def __init__(self, mlflow_client, run_importer=None, await_creation_for=None, host=None):
+        super().__init__(mlflow_client, run_importer, await_creation_for=await_creation_for, host=host)
 
     def import_model(self, model_name, input_dir, experiment_name, delete_model=False, verbose=False, sleep_time=30):
         """
         :param model_name: Model name.
         :param input_dir: Input directory.
-        :param experiment_name: The name of the experiment
+        :param experiment_name: The name of the experiment.
         :param delete_model: Delete current model before importing versions.
         :param verbose: Verbose.
         :param sleep_time: Seconds to wait for model version crreation.
         :return: Model import manifest.
         """
         model_dct = self._import_model(model_name, input_dir, delete_model, verbose, sleep_time)
-        mlflow.set_experiment(experiment_name)
         print("Importing versions:")
+        imported_run_ids = []
         for vr in model_dct["latest_versions"]:
             run_id = self._import_run(input_dir, experiment_name, vr)
+            imported_run_ids.append(run_id)
             self.import_version(model_name, vr, run_id, sleep_time)
         if verbose:
             model_utils.dump_model_versions(self.mlflow_client, model_name)
+        return imported_run_ids
 
     def _import_run(self, input_dir, experiment_name, vr):
         run_id = vr["run_id"]
@@ -149,7 +152,6 @@ def import_model(self, model_name, input_dir, delete_model=False, verbose=False,
         for vr in model_dct["latest_versions"]:
             src_run_id = vr["run_id"]
             dst_run_id = self.run_info_map[src_run_id].run_id
-            mlflow.set_experiment(vr["_experiment_name"])
             self.import_version(model_name, vr, dst_run_id, sleep_time)
         if verbose:
             model_utils.dump_model_versions(self.mlflow_client, model_name)

mlflow_export_import/run/export_run.py

@@ -1,4 +1,4 @@
-""" 
+"""
 Exports a run to a directory.
 """
 
@@ -14,20 +14,19 @@
 from mlflow_export_import.common import MlflowExportImportException
 from mlflow_export_import import utils, click_doc
 
-print("MLflow Version:", mlflow.version.VERSION)
-print("MLflow Tracking URI:", mlflow.get_tracking_uri())
 
 class RunExporter:
-    def __init__(self, mlflow_client, export_source_tags=False, notebook_formats=None):
+    def __init__(self, mlflow_client, export_source_tags=False, notebook_formats=None, host=None):
         """
         :param mlflow_client: MLflow client.
         :param export_source_tags: Export source run metadata tags.
         :param notebook_formats: List of notebook formats to export. Values are SOURCE, HTML, JUPYTER or DBC.
+        :param host: Pass host to the DatabricksHttpClient.
         """
         if notebook_formats is None:
             notebook_formats = []
         self.mlflow_client = mlflow_client
-        self.dbx_client = DatabricksHttpClient()
+        self.dbx_client = DatabricksHttpClient(host=host)
         print("Databricks REST client:", self.dbx_client)
         self.export_source_tags = export_source_tags
         self.notebook_formats = notebook_formats
@@ -38,7 +37,7 @@ def _get_metrics_with_steps(self, run):
             metric_history = self.mlflow_client.get_metric_history(run.info.run_id,metric)
             lst = [utils.strip_underscores(m) for m in metric_history]
             for x in lst:
-                del x["key"] 
+                del x["key"]
             metrics_with_steps[metric] = lst
         return metrics_with_steps
 
@@ -91,8 +90,8 @@ def _export_notebook(self, output_dir, notebook, tags, fs):
         revision_id = tags["mlflow.databricks.notebookRevisionID"]
         notebook_path = tags["mlflow.databricks.notebookPath"]
         notebook_name = os.path.basename(notebook_path)
-        manifest = { 
-           "mlflow.databricks.notebookRevisionID": revision_id, 
+        manifest = {
+           "mlflow.databricks.notebookRevisionID": revision_id,
            "mlflow.databricks.notebookPath": notebook_path,
            "mlflow.databricks.export-notebook-revision": revision_id }
         path = os.path.join(notebook_dir, "manifest.json")
@@ -101,11 +100,11 @@ def _export_notebook(self, output_dir, notebook, tags, fs):
             self._export_notebook_format(notebook_dir, notebook, format, format.lower(), notebook_name, revision_id)
 
     def _export_notebook_format(self, notebook_dir, notebook, format, extension, notebook_name, revision_id):
-        params = { 
-            "path": notebook, 
+        params = {
+            "path": notebook,
             "direct_download": True,
             "format": format,
-            "revision_timestamp": revision_id 
+            "revision_timestamp": revision_id
         }
         try:
             rsp = self.dbx_client._get("workspace/export", params)
@@ -115,26 +114,26 @@ def _export_notebook_format(self, notebook_dir, notebook, format, extension, not
             print(f"WARNING: Cannot save notebook '{notebook}'. {e}")
 
 @click.command()
-@click.option("--run-id", 
-    help="Run ID.", 
+@click.option("--run-id",
+    help="Run ID.",
     type=str,
     required=True
 )
-@click.option("--output-dir", 
-    help="Output directory.", 
+@click.option("--output-dir",
+    help="Output directory.",
     type=str,
     required=True
 )
-@click.option("--export-source-tags", 
-    help=click_doc.export_source_tags, 
-    type=bool, 
-    default=False, 
+@click.option("--export-source-tags",
+    help=click_doc.export_source_tags,
+    type=bool,
+    default=False,
     show_default=True
 )
-@click.option("--notebook-formats", 
-    help=click_doc.notebook_formats, 
+@click.option("--notebook-formats",
+    help=click_doc.notebook_formats,
     type=str,
-    default="", 
+    default="",
     show_default=True
 )
 
@@ -145,7 +144,7 @@ def main(run_id, output_dir, export_source_tags, notebook_formats):
     client = mlflow.tracking.MlflowClient()
     exporter = RunExporter(
       client,
-      export_source_tags=export_source_tags, 
+      export_source_tags=export_source_tags,
       notebook_formats=utils.string_to_list(notebook_formats))
     exporter.export_run(run_id, output_dir)