Move models from one AML workspace to another

CHANGELOG.md

@@ -12,7 +12,7 @@ created.
 ## Upcoming
 
 ### Added
-
+- ([#441](https://github.com/microsoft/InnerEye-DeepLearning/pull/441)) Add script to move models from one AzureML workspace to another: `python InnerEye/Scripts/move_model.py`
 - ([#417](https://github.com/microsoft/InnerEye-DeepLearning/pull/417)) Added a generic way of adding PyTorch Lightning
 models to the toolbox. It is now possible to train almost any Lightning model with the InnerEye toolbox in AzureML,
 with only minimum code changes required. See [the MD documentation](docs/bring_your_own_model.md) for details.

InnerEye/Scripts/move_model.py

@@ -0,0 +1,131 @@
+#  ------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  ------------------------------------------------------------------------------------------
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Tuple
+
+import json
+from attr import dataclass
+from azureml.core import Environment, Model, Workspace
+
+PYTHON_ENVIRONMENT_NAME = "python_environment_name"
+MODEL_PATH = "MODEL"
+ENVIRONMENT_PATH = "ENVIRONMENT"
+MODEL_JSON = "model.json"
+
+
+@dataclass
+class MoveModelConfig:
+    model_id: str
+    path: str
+    action: str
+    workspace_name: str = ""
+    subscription_id: str = ""
+    resource_group: str = ""
+
+    def get_paths(self) -> Tuple[Path, Path]:
+        """
+        Gets paths and creates folders if necessary
+        :param path: Base path
+        :param model_id: The model ID
+        :return: model_path, environment_path
+        """
+        model_id_path = Path(self.path) / self.model_id.replace(":", "_")
+        model_id_path.mkdir(parents=True, exist_ok=True)
+        model_path = model_id_path / MODEL_PATH
+        model_path.mkdir(parents=True, exist_ok=True)
+        env_path = model_id_path / ENVIRONMENT_PATH
+        env_path.mkdir(parents=True, exist_ok=True)
+        return model_path, env_path
+
+
+def download_model(ws: Workspace, config: MoveModelConfig) -> Model:
+    """
+    Downloads an InnerEye model from an AzureML workspace
+    :param ws: The AzureML workspace
+    :param config: move config
+    :return: the exported Model
+    """
+    model = Model(ws, id=config.model_id)
+    model_path, environment_path = config.get_paths()
+    with open(model_path / MODEL_JSON, 'w') as f:
+        json.dump(model.serialize(), f)
+    model.download(target_dir=str(model_path))
+    env_name = model.tags.get(PYTHON_ENVIRONMENT_NAME)
+    environment = ws.environments.get(env_name)
+    environment.save_to_directory(str(environment_path), overwrite=True)
+    return model
+
+
+def upload_model(ws: Workspace, config: MoveModelConfig) -> Model:
+    """
+    Uploads an InnerEye model to an AzureML workspace
+    :param ws: The AzureML workspace
+    :param config: move config
+    :return: imported Model
+    """
+    model_path, environment_path = config.get_paths()
+    with open(model_path / MODEL_JSON, 'r') as f:
+        model_dict = json.load(f)
+
+    new_model = Model.register(ws, model_path=str(model_path / "final_model"), model_name=model_dict['name'],
+                               tags=model_dict['tags'], properties=model_dict['properties'],
+                               description=model_dict['description'])
+    env = Environment.load_from_directory(str(environment_path))
+    env.register(workspace=ws)
+    print(f"Environment {env.name} registered")
+    return new_model
+
+
+def get_workspace(config: MoveModelConfig) -> Workspace:
+    """
+    Get workspace based on command line input config
+    :param config: MoveModelConfig
+    :return: an Azure ML workspace
+    """
+    return Workspace.get(name=config.workspace_name, subscription_id=config.subscription_id,
+                         resource_group=config.resource_group)
+
+
+def main() -> None:
+    parser = ArgumentParser()
+    parser.add_argument("-a", "--action", type=str, required=True,
+                        help="Action (download or upload)")
+    parser.add_argument("-w", "--workspace_name", type=str, required=True,
+                        help="Azure ML workspace name")
+    parser.add_argument("-s", "--subscription_id", type=str, required=True,
+                        help="AzureML subscription id")
+    parser.add_argument("-r", "--resource_group", type=str, required=True,
+                        help="AzureML resource group")
+    parser.add_argument("-p", "--path", type=str, required=True,
+                        help="The path to download or upload model")
+    parser.add_argument("-m", "--model_id", type=str, required=True,
+                        help="The AzureML model ID")
+
+    args = parser.parse_args()
+    config = MoveModelConfig(workspace_name=args.workspace_name, subscription_id=args.subscription_id,
+                             resource_group=args.resource_group,
+                             path=args.path, action=args.action, model_id=args.model_id)
+    ws = get_workspace(config)
+    move(ws, config)
+
+
+def move(ws: Workspace, config: MoveModelConfig) -> Model:
+    """
+    Moves a model: downloads or uploads the model depending on the configs
+    :param config: the move model config
+    :param ws: The Azure ML workspace
+    :return: the download or upload model
+    """
+    if config.action == "download":
+        return download_model(ws, config)
+    elif config.action == "upload":
+        return upload_model(ws, config)
+    else:
+        raise ValueError(f'Invalid action {config.action}, allowed values: import or export')
+
+
+if __name__ == "__main__":
+    main()

README.md

@@ -99,6 +99,7 @@ Further detailed instructions, including setup in Azure, are here:
 1. [Sample Segmentation and Classification tasks](docs/sample_tasks.md)
 1. [Debugging and monitoring models](docs/debugging_and_monitoring.md)
 1. [Model diagnostics](docs/model_diagnostics.md)
+1. [Move a model to a different workspace](docs/move_model.md)   
 1. [Deployment](docs/deploy_on_aml.md)
 
 ![docs/deployment.png](docs/deployment.png)

Tests/Scripts/test_move_model.py

@@ -0,0 +1,28 @@
+#  ------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  ------------------------------------------------------------------------------------------
+
+from InnerEye.Azure.azure_config import AzureConfig
+from InnerEye.Common import fixed_paths
+from InnerEye.Common.output_directories import OutputFolderForTests
+from InnerEye.Scripts.move_model import MoveModelConfig, PYTHON_ENVIRONMENT_NAME, move
+
+MODEL_ID = "PassThroughModel:1"
+
+
+def test_download_and_upload(test_output_dirs: OutputFolderForTests) -> None:
+    """
+    Test that downloads and uploads a model to a workspace
+    """
+    azure_config = AzureConfig.from_yaml(yaml_file_path=fixed_paths.SETTINGS_YAML_FILE,
+                                         project_root=fixed_paths.repository_root_directory())
+    ws = azure_config.get_workspace()
+    config_download = MoveModelConfig(model_id=MODEL_ID, path=str(test_output_dirs.root_dir), action="download")
+    move(ws, config_download)
+    assert (test_output_dirs.root_dir / MODEL_ID.replace(":", "_")).is_dir()
+    config_upload = MoveModelConfig(model_id=MODEL_ID, path=str(test_output_dirs.root_dir), action="upload")
+    model = move(ws, config_upload)
+    assert model is not None
+    assert PYTHON_ENVIRONMENT_NAME in model.tags
+    assert model.description != ""

docs/move_model.md

@@ -0,0 +1,19 @@
+# Move a model to other workspace
+
+The InnerEye models on AzureML are composed of two parts in order to reproduce the same settings at inference and
+training time:
+
+- Model: The model is registered in the AzureML registry and contains the code used at training time and pytorch
+  checkpoint
+- Environment: The Azure ML environment used to train the model. This contains the docker image with all the
+  dependencies that were used for training
+
+If you want to export a model from one Workspace to another you can use the following command to download and upload a model
+from an AzureML workspace. This script does not use settings.yml, it uses interactive authentication, and the workspace specified in the
+parameters. The model will be written to the path in --path parameter with two folders one for the `MODEL` and one for the `ENVIRONMENT` files.
+
+- Download to
+  path: `python InnerEye/Scripts/move_model.py -a download --path ./ --workspace_name "<name>" --resource_group "<name>" --subscription_id "<sub_id>" --model_id "name:version"`
+
+- Upload from
+  path: `python InnerEye/Scripts/move_model.py - upload --path ./ --workspace_name "<name>" --resource_group "<name>" --subscription_id "<sub_id>" --model_id "name:version"`