Upgrade model management to SDK v2

This article gives a comparison of scenario(s) in SDK v1 and SDK v2.

Create model

  • SDK v1

    import urllib.request
    from azureml.core.model import Model
    
    # Register model
    model = Model.register(ws, model_name="local-file-example", model_path="mlflow-model/model.pkl")
    
  • SDK v2

    from azure.ai.ml.entities import Model
    from azure.ai.ml.constants import AssetTypes
    
    file_model = Model(
        path="mlflow-model/model.pkl",
        type=AssetTypes.CUSTOM_MODEL,
        name="local-file-example",
        description="Model created from local file."
    )
    ml_client.models.create_or_update(file_model)
    

Use model in an experiment/job

  • SDK v1

    model = run.register_model(model_name='run-model-example',
                               model_path='outputs/model/')
    print(model.name, model.id, model.version, sep='\t')
    
  • SDK v2

    from azure.ai.ml.entities import Model
    from azure.ai.ml.constants import AssetTypes
    
    run_model = Model(
        path="azureml://jobs/$RUN_ID/outputs/artifacts/paths/model/",
        name="run-model-example",
        description="Model created from run.",
        type=AssetTypes.CUSTOM_MODEL
    )
    
    ml_client.models.create_or_update(run_model)
    

For more information about models, see Work with models in Azure Machine Learning.

Mapping of key functionality in SDK v1 and SDK v2

Next steps

For more information, see the documentation here: