mlflow / mlflow-export-import

Apache License 2.0
132 stars 78 forks source link

Add versions to export import of models, resolve global variables, pass host around the initializers #23

Open kongzii opened 2 years ago

kongzii commented 2 years ago

Hi, thank you for this project. I am working on a simple script that will use this to migrate runs and models between our development and production mlflow server in the company, however, I found a couple of issues so I tried to fix them.

The result is that now, you are able to use this library from the other python scripts, without need of covering environment variables, example:

       MODEL_NAME = '...'
       EXPERIMENT_NAME = '...'
       TMP_DIR = '...'
       SOURCE_TRACKING_URI = '...'
       TARGET_TRACKING_URI = '...'

        export_client = mlflow.tracking.MlflowClient(tracking_uri=SOURCE_TRACKING_URI)
        model_exporter = export_model.ModelExporter(
            mlflow_client=export_client,
            host=SOURCE_TRACKING_URI,
            versions=["11", "12"],
        )
        model_exporter.export_model(MODEL_NAME, TMP_DIR)

        import_client = mlflow.tracking.MlflowClient(tracking_uri=TARGET_TRACKING_URI)
        model_importer = import_model.ModelImporter(
            mlflow_client=import_client,
            host=TARGET_TRACKING_URI,
            await_creation_for=600,
        )
        exported_run_id = model_importer.import_model(
            MODEL_NAME,
            TMP_DIR,
            experiment_name=EXPERIMENT_NAME,
        )

I will do probably more contributions in the future, but I wanted to get your feedback before there is too many of changes.

kongzii commented 2 years ago

Hi @amesar, sorry to bother you, but do you have any feedback for this PR, please? Is there something I should change?

kongzii commented 2 years ago

Hi @amesar, any feedback for this, please? If you don't want to merge this, no problem, I can fork the project and do the required modifications just for me, but I thought it would be a shame to not collaborate on a project like this.