def __init__(self, client=None, export_metadata_tags=False, notebook_formats=["SOURCE"], filesystem=None): self.client = client or mlflow.tracking.MlflowClient() self.dbx_client = DatabricksHttpClient() print("Databricks REST client:",self.dbx_client) self.fs = filesystem or _filesystem.get_filesystem() print("Filesystem:",type(self.fs).__name__) self.export_metadata_tags = export_metadata_tags self.notebook_formats = notebook_formats
def __init__(self, client=None, export_metadata_tags=False, notebook_formats=["SOURCE"], filesystem=None): self.client = client or mlflow.tracking.MlflowClient() self.fs = filesystem or _filesystem.get_filesystem() print("Filesystem:", type(self.fs).__name__) self.run_exporter = RunExporter(self.client, export_metadata_tags, notebook_formats, self.fs)
def __init__(self, export_metadata_tags=False, notebook_formats=["SOURCE"], filesystem=None): self.fs = filesystem or _filesystem.get_filesystem() self.client = mlflow.tracking.MlflowClient() self.client2 = HttpClient("api/2.0/preview/mlflow") self.run_exporter = RunExporter( self.client, export_metadata_tags=export_metadata_tags, notebook_formats=notebook_formats, filesystem=filesystem)
def _export_models(models, output_dir, notebook_formats, export_notebook_revision, stages, export_run=True, use_threads=False): max_workers = os.cpu_count() or 4 if use_threads else 1 start_time = time.time() if models == "all": models = [ model.name for model in client.list_registered_models() ] elif models.endswith("*"): model_prefix = models[:-1] models = [ model.name for model in client.list_registered_models() if model.name.startswith(model_prefix) ] # Wish there was an model search method for efficiency] else: models = models.split(",") print("Models:") for model in models: print(f" {model}") exporter = ModelExporter(stages=stages, notebook_formats=utils.string_to_list(notebook_formats), export_notebook_revision=export_notebook_revision, export_run=export_run) futures = [] with ThreadPoolExecutor(max_workers=max_workers) as executor: for model in models: dir = os.path.join(output_dir, model) future = executor.submit(exporter.export_model, model, dir) futures.append(future) ok_models = [] ; failed_models = [] for future in futures: result = future.result() if result[0]: ok_models.append(result[1]) else: failed_models.append(result[1]) duration = round(time.time() - start_time, 1) manifest = { "info": { "mlflow_version": mlflow.__version__, "mlflow_tracking_uri": mlflow.get_tracking_uri(), "export_time": utils.get_now_nice(), "total_models": len(models), "ok_models": len(ok_models), "failed_models": len(failed_models), "duration": duration }, "stages": stages, "notebook_formats": notebook_formats, "export_notebook_revision": export_notebook_revision, "ok_models": ok_models, "failed_models": failed_models } fs = _filesystem.get_filesystem(output_dir) fs.mkdirs(output_dir) with open(os.path.join(output_dir, "manifest.json"), "w") as f: f.write(json.dumps(manifest, indent=2)+"\n") print(f"{len(models)} models exported") print(f"Duration for registered models export: {duration} seconds")
def _export_model(self, model_name, output_dir): fs = _filesystem.get_filesystem(output_dir) model = self.http_client.get(f"registered-models/get", {"name": model_name}) fs.mkdirs(output_dir) model["registered_model"]["latest_versions"] = [] versions = self.mlflow_client.search_model_versions( f"name='{model_name}'") print(f"Found {len(versions)} versions for model {model_name}") manifest = [] exported_versions = 0 for vr in versions: if len(self.stages) > 0 and not vr.current_stage.lower( ) in self.stages: continue run_id = vr.run_id opath = os.path.join(output_dir, run_id) opath = opath.replace("dbfs:", "/dbfs") dct = { "version": vr.version, "stage": vr.current_stage, "run_id": run_id } print(f"Exporting: {dct}") manifest.append(dct) try: if self.export_run: self.run_exporter.export_run(run_id, opath) 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) dct["_experiment_name"] = experiment.name model["registered_model"]["latest_versions"].append(dct) exported_versions += 1 except mlflow.exceptions.RestException as e: if "RESOURCE_DOES_NOT_EXIST: Run" in str(e): print( f"WARNING: Run for version {vr.version} does not exist. {e}" ) else: import traceback traceback.print_exc() print( f"Exported {exported_versions}/{len(versions)} versions for model {model_name}" ) path = os.path.join(output_dir, "model.json") utils.write_json_file(fs, path, model) return manifest
def export_experiment(self, exp_id_or_name, output_dir, run_ids=None): """ :param exp_id_or_name: Experiment ID or name. :param output_dir: Output directory. :param run_ids: List of run IDs to export. If None export all run IDs. :return: Number of successful and number of failed runs. """ exp = mlflow_utils.get_experiment(self.mlflow_client, exp_id_or_name) exp_id = exp.experiment_id print( f"Exporting experiment '{exp.name}' (ID {exp.experiment_id}) to '{output_dir}'" ) fs = _filesystem.get_filesystem(output_dir) print("Filesystem:", type(fs).__name__) fs.mkdirs(output_dir) exp = self.mlflow_client.get_experiment(exp_id) dct = {"experiment": utils.strip_underscores(exp)} ok_run_ids = [] failed_run_ids = [] j = -1 if run_ids: for j, run_id in enumerate(run_ids): run = self.mlflow_client.get_run(run_id) self._export_run(j, run, output_dir, ok_run_ids, failed_run_ids) else: for j, run in enumerate( SearchRunsIterator(self.mlflow_client, exp_id)): self._export_run(j, run, output_dir, ok_run_ids, failed_run_ids) dct["export_info"] = { "export_time": utils.get_now_nice(), "num_total_runs": (j + 1), "num_ok_runs": len(ok_run_ids), "ok_runs": ok_run_ids, "num_failed_runs": len(failed_run_ids), "failed_runs": failed_run_ids } path = os.path.join(output_dir, "manifest.json") utils.write_json_file(fs, path, dct) msg = f"for experiment '{exp.name}' (ID: {exp.experiment_id})" if len(failed_run_ids) == 0: print(f"All {len(ok_run_ids)} runs succesfully exported {msg}") else: print(f"{len(ok_run_ids)/j} runs succesfully exported {msg}") print(f"{len(failed_run_ids)/j} runs failed {msg}") return len(ok_run_ids), len(failed_run_ids)
def import_all(input_dir, delete_model, use_src_user_id, import_mlflow_tags, import_metadata_tags, verbose, use_threads): start_time = time.time() exp_res = import_experiments(input_dir, use_src_user_id, import_mlflow_tags, import_metadata_tags) run_info_map = _remap(exp_res[0]) model_res = import_models(input_dir, run_info_map, delete_model, verbose, use_threads) duration = round(time.time() - start_time, 1) dct = { "duration": duration, "experiment_import": exp_res[1], "model_import": model_res } fs = _filesystem.get_filesystem(".") utils.write_json_file(fs, "import_report.json", dct) print("\nImport report:") print(json.dumps(dct, indent=2) + "\n")
def export_run(self, run_id, output_dir): """ :param run_id: Run ID. :param output_dir: Output directory. :return: whether export succeeded. """ fs = _filesystem.get_filesystem(output_dir) run = self.mlflow_client.get_run(run_id) fs.mkdirs(output_dir) tags = utils.create_tags_for_metadata(self.mlflow_client, run, self.export_metadata_tags) dct = { "info": utils.strip_underscores(run.info) , "params": run.data.params, "metrics": run.data.metrics, "tags": tags, } path = os.path.join(output_dir,"run.json") utils.write_json_file(fs, path, dct) # copy artifacts dst_path = os.path.join(output_dir,"artifacts") try: TAG_NOTEBOOK_PATH = "mlflow.databricks.notebookPath" artifacts = self.mlflow_client.list_artifacts(run.info.run_id) if len(artifacts) > 0: # Because of https://github.com/mlflow/mlflow/issues/2839 fs.mkdirs(dst_path) self.mlflow_client.download_artifacts(run.info.run_id,"", dst_path=mk_local_path(dst_path)) notebook = tags.get(TAG_NOTEBOOK_PATH, None) if notebook is not None: if len(self.notebook_formats) > 0: self.export_notebook(output_dir, notebook, run.data.tags, fs) elif len(self.notebook_formats) > 0: print(f"WARNING: Cannot export notebook since tag '{TAG_NOTEBOOK_PATH}' is not set.") return True except Exception as e: print("ERROR: run_id:",run.info.run_id,"Exception:",e) traceback.print_exc() return False
def export_experiments(experiments, output_dir, export_metadata_tags, notebook_formats, export_notebook_revision=False, use_threads=False): """ :param: experiments: Can be either: - List of experiment names - List of experiment IDs - Dictionary with experiment ID key and list of run IDs - String with comma-delimited experiment names or IDs. """ start_time = time.time() max_workers = os.cpu_count() or 4 if use_threads else 1 export_all_runs = not isinstance(experiments, dict) if export_all_runs: experiments = utils.get_experiments(experiments) table_data = experiments columns = ["Experiment ID"] experiments_dct = {} else: experiments_dct = experiments experiments = experiments.keys() experiments = utils.get_experiments(experiments) table_data = [[exp_id, len(runs)] for exp_id, runs in experiments_dct.items()] num_runs = sum(x[1] for x in table_data) table_data.append(["Total", num_runs]) columns = ["Experiment ID", "# Runs"] utils.show_table("Experiments", table_data, columns) print("") ok_runs = 0 failed_runs = 0 export_results = [] futures = [] exporter = ExperimentExporter(client, export_metadata_tags, utils.string_to_list(notebook_formats), export_notebook_revision) with ThreadPoolExecutor(max_workers=max_workers) as executor: for exp_id_or_name in experiments: run_ids = experiments_dct.get(exp_id_or_name, None) future = executor.submit(_export_experiment, exp_id_or_name, output_dir, exporter, export_results, run_ids) futures.append(future) duration = round(time.time() - start_time, 1) ok_runs = 0 failed_runs = 0 for future in futures: result = future.result() ok_runs += result[0] failed_runs += result[1] total_runs = ok_runs + failed_runs duration = round(time.time() - start_time, 1) dct = { "info": { "mlflow_version": mlflow.__version__, "mlflow_tracking_uri": mlflow.get_tracking_uri(), "export_time": utils.get_now_nice(), "duration": duration, "experiments": len(experiments), "total_runs": total_runs, "ok_runs": ok_runs, "failed_runs": failed_runs }, "experiments": export_results } fs = _filesystem.get_filesystem(output_dir) fs.mkdirs(output_dir) with open(os.path.join(output_dir, "manifest.json"), "w") as f: f.write(json.dumps(dct, indent=2) + "\n") print(f"{len(experiments)} experiments exported") print(f"{ok_runs}/{total_runs} runs succesfully exported") if failed_runs > 0: print(f"{failed_runs}/{total_runs} runs failed") print(f"Duration for experiments export: {duration} seconds")
def __init__(self, filesystem=None): self.fs = filesystem or _filesystem.get_filesystem() self.client = mlflow.tracking.MlflowClient() self.run_importer = RunImporter(self.client, mlmodel_fix=True)
def __init__(self, filesystem=None, run_importer=None): self.fs = filesystem or _filesystem.get_filesystem() self.client = mlflow.tracking.MlflowClient() self.run_importer = run_importer if run_importer else RunImporter( self.client, mlmodel_fix=True, import_mlflow_tags=True)