Versions

Signed-off-by: Peter Jung <peter@jung.ninja>

Author: kongzii

.gitignore

@@ -0,0 +1,3 @@
+# Test files
+info.txt
+run_tests.log

mlflow_export_import/click_doc.py

@@ -9,6 +9,8 @@
 
 model_stages = "Stages to export (comma seperated). Default is all stages. Values are Production, Staging, Archived and None."
 
+model_versions = "Versions to export (comma seperated). Default is all versions. Values are valid integer numbers."
+
 delete_model = "If the model exists, first delete the model and all its versions."
 
 use_threads = "Process export/import in parallel using threads."

mlflow_export_import/model/export_model.py

@@ -11,18 +11,20 @@
 from mlflow_export_import import utils, click_doc
 
 class ModelExporter():
-    def __init__(self,  mlflow_client, export_source_tags=False, notebook_formats=None, stages=None, export_run=True):
+    def __init__(self,  mlflow_client, export_source_tags=False, notebook_formats=None, stages=None, versions=None, export_run=True):
         """
         :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.
         """
         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.stages = self._normalize_stages(stages)
+        self.versions = self._normalize_versions(versions)
         self.export_run = export_run
 
     def export_model(self, model_name, output_dir):
@@ -50,6 +52,8 @@ def _export_model(self, model_name, output_dir):
         for vr in versions:
             if len(self.stages) > 0 and not vr.current_stage.lower() in self.stages:
                 continue
+            if len(self.versions) > 0 and vr.version not in self.versions:
+                continue
             run_id = vr.run_id
             opath = os.path.join(output_dir,run_id)
             opath = opath.replace("dbfs:", "/dbfs")
@@ -89,14 +93,26 @@ def _normalize_stages(self, stages):
                 print(f"WARNING: stage '{stage}' must be one of: {model_version_stages.ALL_STAGES}")
         return stages
 
+    def _normalize_versions(self, versions):
+        if versions is None:
+            return []
+        if isinstance(versions, str):
+            versions = versions.split(",")
+        for version in versions:
+            try:
+                int(version)
+            except ValueError:
+                print(f"WARNING: version '{version}' must be a valid number")
+        return versions
+
 @click.command()
-@click.option("--model", 
-    help="Registered model name.", 
+@click.option("--model",
+    help="Registered model name.",
     type=str,
     required=True
 )
-@click.option("--output-dir", 
-    help="Output directory.", 
+@click.option("--output-dir",
+    help="Output directory.",
     type=str,
     required=True
 )
@@ -106,24 +122,29 @@ def _normalize_stages(self, stages):
     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
 )
-@click.option("--stages", 
-    help=click_doc.model_stages, 
+@click.option("--stages",
+    help=click_doc.model_stages,
+    type=str,
+    required=False
+)
+@click.option("--versions",
+    help=click_doc.model_versions,
     type=str,
     required=False
 )
 
-def main(model, output_dir, export_source_tags, notebook_formats, stages):
+def main(model, output_dir, export_source_tags, notebook_formats, stages, versions):
     print("Options:")
     for k,v in locals().items():
         print(f"  {k}: {v}")
     client = mlflow.tracking.MlflowClient()
-    exporter = ModelExporter(client, export_source_tags=export_source_tags, notebook_formats=utils.string_to_list(notebook_formats), stages=stages)
+    exporter = ModelExporter(client, export_source_tags=export_source_tags, notebook_formats=utils.string_to_list(notebook_formats), stages=stages, versions=versions)
     exporter.export_model(model, output_dir)
 
 if __name__ == "__main__":

mlflow_export_import/model/import_model.py

@@ -19,9 +19,9 @@ def __init__(self, mlflow_client, run_importer=None, await_creation_for=None):
         :param run_importer: RunImporter instance.
         :param await_creation_for: Seconds to wait for model version crreation.
         """
-        self.mlflow_client = mlflow_client 
+        self.mlflow_client = mlflow_client
         self.run_importer = run_importer if run_importer else RunImporter(self.mlflow_client, mlmodel_fix=True)
-        self.await_creation_for = await_creation_for 
+        self.await_creation_for = await_creation_for
 
     def _import_version(self, model_name, src_vr, dst_run_id, dst_source, sleep_time):
         """
@@ -33,10 +33,10 @@ def _import_version(self, model_name, src_vr, dst_run_id, dst_source, sleep_time
         """
         src_current_stage = src_vr["current_stage"]
         dst_source = dst_source.replace("file://","") # OSS MLflow
-        if not dst_source.startswith("dbfs:") and not os.path.exists(dst_source):
+        if not dst_source.startswith("dbfs:") and not dst_source.startswith("s3:") and not os.path.exists(dst_source):
             raise MlflowExportImportException(f"'source' argument for MLflowClient.create_model_version does not exist: {dst_source}")
         kwargs = {"await_creation_for": self.await_creation_for } if self.await_creation_for else {}
-        version = self.mlflow_client.create_model_version(model_name, dst_source, dst_run_id, **kwargs)
+        version = self.mlflow_client.create_model_version(model_name, dst_source, dst_run_id, src_vr["tags"], **kwargs)
         model_utils.wait_until_version_is_ready(self.mlflow_client, model_name, version, sleep_time=sleep_time)
         if src_current_stage != "None":
             self.mlflow_client.transition_model_version_stage(model_name, version.version, src_current_stage)
@@ -173,46 +173,46 @@ def _path_join(x,y):
     """ Account for DOS backslash """
     path = os.path.join(x,y)
     if path.startswith("dbfs:"):
-        path = path.replace("\\","/") 
+        path = path.replace("\\","/")
     return path
 
 @click.command()
-@click.option("--input-dir", 
-    help="Input directory produced by export_model.py.", 
+@click.option("--input-dir",
+    help="Input directory produced by export_model.py.",
     type=str,
     required=True
 )
-@click.option("--model", 
-    help="New registered model name.", 
+@click.option("--model",
+    help="New registered model name.",
     type=str,
-    required=True, 
+    required=True,
 )
-@click.option("--experiment-name", 
-    help="Destination experiment name  - will be created if it does not exist.", 
+@click.option("--experiment-name",
+    help="Destination experiment name  - will be created if it does not exist.",
     type=str,
     required=True
 )
-@click.option("--delete-model", 
-    help=click_doc.delete_model, 
+@click.option("--delete-model",
+    help=click_doc.delete_model,
     type=bool,
-    default=False, 
+    default=False,
     show_default=True
 )
-@click.option("--await-creation-for", 
-    help="Await creation for specified seconds.", 
-    type=int, 
-    default=None, 
+@click.option("--await-creation-for",
+    help="Await creation for specified seconds.",
+    type=int,
+    default=None,
     show_default=True
 )
-@click.option("--sleep-time", 
-    help="Sleep time for polling until version.status==READY.", 
+@click.option("--sleep-time",
+    help="Sleep time for polling until version.status==READY.",
     type=int,
     default=5,
 )
-@click.option("--verbose", 
-    help="Verbose.", 
-    type=bool, 
-    default=False, 
+@click.option("--verbose",
+    help="Verbose.",
+    type=bool,
+    default=False,
     show_default=True
 )
 

tests/run_tests.sh

@@ -27,13 +27,13 @@ export MLFLOW_TRACKING_URI_SRC=http://localhost:${PORT_SRC}
 export MLFLOW_TRACKING_URI_DST=http://localhost:${PORT_DST}
 
 message() {
-  echo 
+  echo
   echo "******************************************************"
   echo "*"
   echo "* $*"
   echo "*"
   echo "******************************************************"
-  echo 
+  echo
 }
 
 run_tests() {
@@ -50,7 +50,7 @@ launch_server() {
   mlflow server \
     --host localhost --port ${port}  \
     --backend-store-uri sqlite:///mlflow_${port}.db \
-    --default-artifact-root $PWD/mlruns_${port}
+    --default-artifact-root "$(PWD)/mlruns_${port}"
 }
 
 kill_server() {

tests/test_models.py

@@ -1,7 +1,7 @@
 from mlflow_export_import.model.export_model import ModelExporter
 from mlflow_export_import.model.import_model import ModelImporter
 from mlflow_export_import import utils
-import utils_test 
+import utils_test
 import compare_utils
 from init_tests import mlflow_context
 
@@ -57,6 +57,34 @@ def test_export_import_model_stages(mlflow_context):
         [vr_prod_src, vr_staging_src],
         [vr_prod_dst, vr_staging_dst])
 
+def test_export_import_model_versions(mlflow_context):
+    model_name_src = utils_test.mk_test_object_name_default()
+    model_src = mlflow_context.client_src.create_registered_model(model_name_src)
+
+    vr_staging_src = _create_version(mlflow_context.client_src, model_name_src, "Staging")
+    vr_prod_src = _create_version(mlflow_context.client_src, model_name_src, "Production")
+
+    exporter = ModelExporter(mlflow_context.client_src, versions=[vr_staging_src.version, vr_prod_src.version])
+    exporter.export_model(model_name_src, mlflow_context.output_dir)
+
+    model_name_dst = utils_test.create_dst_model_name(model_name_src)
+    experiment_name = model_name_dst
+    importer = ModelImporter(mlflow_context.client_dst)
+    importer.import_model(model_name_dst, mlflow_context.output_dir, experiment_name, delete_model=True, verbose=False, sleep_time=10)
+
+    model_dst = mlflow_context.client_dst.get_registered_model(model_name_dst)
+    assert len(model_dst.latest_versions) == 2
+
+    versions = mlflow_context.client_dst.get_latest_versions(model_name_dst)
+    vr_prod_dst = [vr for vr in versions if vr.version == vr_prod_src.version][0]
+    vr_staging_dst = [vr for vr in versions if vr.version == vr_staging_src.version][0]
+
+    _compare_models(model_src, model_dst)
+    _compare_version_lists(
+        mlflow_context, mlflow_context.output_dir,
+        [vr_prod_src, vr_staging_src],
+        [vr_prod_dst, vr_staging_dst])
+
 
 def _create_version(client, model_name, stage=None):
     run = _create_run(client)
@@ -69,7 +97,7 @@ def _create_version(client, model_name, stage=None):
 def _create_run(client):
     _, run = utils_test.create_simple_run(client)
     return client.get_run(run.info.run_id)
-    
+
 def _compare_models(model_src, model_dst):
     assert model_src.description == model_dst.description
     assert model_src.tags == model_dst.tags
@@ -85,7 +113,7 @@ def _compare_versions(mlflow_context, output_dir, vr_src, vr_dst):
     assert vr_src.status == vr_dst.status
     assert vr_src.status_message == vr_dst.status_message
     if not utils.importing_into_databricks():
-        assert vr_src.user_id == vr_dst.user_id 
+        assert vr_src.user_id == vr_dst.user_id
     assert vr_src.run_id != vr_dst.run_id
     run_src = mlflow_context.client_src.get_run(vr_src.run_id)
     run_dst = mlflow_context.client_dst.get_run(vr_dst.run_id)