Model management maap v2

assets/training/model_management/components/convert_model_to_mlflow/spec.yaml

@@ -1,15 +1,15 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: convert_model_to_mlflow
-version: 0.0.37
+version: 0.0.38
 type: command
 
 is_deterministic: True
 
 display_name: Convert models to MLflow
 description: Component converts models from supported frameworks to MLflow model packaging format
 
-environment: azureml://registries/azureml/environments/model-management/versions/41
+environment: azureml://registries/azureml/environments/model-management/versions/48
 
 code: ../../src/
 command: |

assets/training/model_management/components/download_model/asset.yaml

@@ -0,0 +1,3 @@
+type: component
+spec: spec.yaml
+categories: ["Models"]

assets/training/model_management/components/download_model/spec.yaml

@@ -0,0 +1,67 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
+
+name: download_model
+version: 0.0.32
+type: command
+
+is_deterministic: True
+
+display_name: Download model
+description: Downloads a publicly available model
+
+environment: azureml://registries/azureml/environments/model-management/versions/48
+
+code: ../../src/
+command: >
+  python run_model_download.py
+  --model-source ${{inputs.model_source}}
+  --model-id '${{inputs.model_id}}'
+  $[[--update-existing-model ${{inputs.update_existing_model}}]]
+  $[[--validation-info ${{inputs.validation_info}}]]
+  $[[--token ${{inputs.token}}]]
+  --model-download-metadata ${{outputs.model_download_metadata}}
+  --model-output-dir ${{outputs.model_output}}
+
+inputs:
+  model_source:
+    type: string
+    description: Storage containers from where model will be sourced from.
+    default: Huggingface
+    enum:
+      - AzureBlob
+      - GIT
+      - Huggingface
+
+  model_id:
+    type: string
+    description: A valid model id for the model source selected. For example you can specify `bert-base-uncased` for importing HuggingFace bert base uncased model. Please specify the complete URL if **GIT** or **AzureBlob** is selected in `model_source`
+
+  validation_info:
+    type: uri_file
+    description: Path to the validation info file
+    optional: true
+
+  update_existing_model:
+    type: boolean
+    default: false
+    description: If set to true, will update the existing model. If set to false, will create a new model.
+    optional: true
+
+  token:
+    type: string
+    description: If set use it to access the private models or authenticate the user. For example, user can get the token for HF private model by creating account in Huggingface, accept the condition for models that needs to be downloaded and create access token from browser. For more details please visit - https://huggingface.co/docs/hub/security-tokens
+    optional: true
+
+outputs:
+  model_download_metadata:
+    type: uri_file
+    description: File name to which model download details will be written. File would contain details that could be useful for model registration in forms of model tags and properties
+
+  model_output:
+    type: uri_folder
+    description: Path to the dowloaded model
+    mode: rw_mount
+
+tags:
+    Preview: ""
+

assets/training/model_management/environments/model-management/context/requirements.txt

@@ -8,8 +8,6 @@ sentencepiece=={{latest-pypi-version}}
 accelerate=={{latest-pypi-version}}
 wget=={{latest-pypi-version}}
 applicationinsights=={{latest-pypi-version}}
-azureml-automl-core=={{latest-pypi-version}}
-azureml-telemetry=={{latest-pypi-version}}
 azureml-automl-dnn-vision=={{latest-pypi-version}}
 pyarrow==14.0.2
 GitPython=={{latest-pypi-version}}

assets/training/model_management/src/azureml/model/mgmt/utils/common_utils.py

@@ -11,16 +11,11 @@
 import time
 from argparse import Namespace
 from azure.ai.ml import MLClient
-from azureml._common._error_definition import AzureMLError
-from azureml._common.exceptions import AzureMLException
+from azure.ai.ml.exceptions import ErrorTarget, ErrorCategory, MlException
 from azure.ai.ml.identity import AzureMLOnBehalfOfCredential
 from azure.identity import ManagedIdentityCredential
 from azureml.core.run import Run
-from azureml.model.mgmt.utils.exceptions import (
-    GenericRunCMDError,
-    HuggingFaceErrorInFetchingModelInfo,
-    UserIdentityMissingError
-    )
+from azureml.model.mgmt.utils.exceptions import ModelImportErrorStrings
 from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
@@ -111,7 +106,12 @@ def get_mlclient(registry_name: str = None):
             # Check if given credential can get token successfully.
             credential.get_token("https://management.azure.com/.default")
         except Exception as ex:
-            raise AzureMLException._with_error(AzureMLError.create(UserIdentityMissingError, exception=ex))
+            message = ModelImportErrorStrings.USER_IDENTITY_MISSING_ERROR
+            raise MlException(
+                message=message, no_personal_data_message=message,
+                error_category=ErrorCategory.USER_ERROR, target=ErrorTarget.COMPONENT,
+                error=ex
+            )
 
     if registry_name is None:
         run = Run.get_context(allow_offline=False)
@@ -280,8 +280,12 @@ def fetch_huggingface_model_info(model_id) -> ModelInfo:
             if model_id == info.modelId:
                 return info
     except Exception as e:
-        raise AzureMLException._with_error(
-            AzureMLError.create(HuggingFaceErrorInFetchingModelInfo, model_id=model_id, error=e)
+        message = ModelImportErrorStrings.ERROR_FETCHING_HUGGING_FACE_MODEL_INFO
+        raise MlException(
+            message=message.format(model_id=model_id, error=""),
+            no_personal_data_message=message.format(model_id=model_id, error=e),
+            error_category=ErrorCategory.USER_ERROR, target=ErrorTarget.COMPONENT,
+            error=e
         )
 
 
@@ -347,7 +351,12 @@ def get_git_lfs_blob_size_in_kb(git_dir: Path) -> int:
         logger.info(f"total size: {stdout} KB")
         return int(stdout)
     except Exception as e:
-        raise AzureMLException._with_error(AzureMLError.create(GenericRunCMDError, error=e))
+        message = ModelImportErrorStrings.CMD_EXECUTION_ERROR
+        raise MlException(
+            message=message, no_personal_data_message=message.format(error=e),
+            error_category=ErrorCategory.SYSTEM_ERROR, target=ErrorTarget.COMPONENT,
+            error=e
+        )
 
 
 class MlflowMetaConstants:

assets/training/model_management/src/azureml/model/mgmt/utils/exceptions.py

@@ -6,15 +6,7 @@
 import time
 import logging
 from functools import wraps
-from azureml._common.exceptions import AzureMLException
-from azureml._common._error_definition.azureml_error import AzureMLError  # type: ignore
-from azureml._common._error_definition.system_error import ClientError  # type: ignore
-from azureml._common._error_definition.user_error import (
-    ArgumentInvalid,
-    Authentication,
-    NotSupported,
-    ConnectionFailure
-)  # type: ignore
+from azure.ai.ml.exceptions import ErrorTarget, ErrorCategory, MlException
 
 from azureml.core.run import Run  # type: ignore
 from azureml.automl.core._run import run_lifecycle_utilities
@@ -58,140 +50,6 @@ class ModelImportErrorStrings:
     )
 
 
-class ModelImportException(AzureMLException):
-    """Base exception for Model Import handling."""
-
-    def __init__(self, exception_message, **kwargs):
-        """Initialize a new instance of LLMException.
-
-        :param exception_message: A message describing the error
-        :type exception_message: str
-        """
-        super(ModelImportException, self).__init__(exception_message, **kwargs)
-
-    @property
-    def error_code(self):
-        """Return error code for azureml_error."""
-        return self._azureml_error.error_definition.code
-
-
-class ModelImportError(ClientError):
-    """Internal Import Model Generic Error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.LOG_UNSAFE_GENERIC_ERROR
-
-
-class GITCloneError(ConnectionFailure):
-    """GIT clone error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.GIT_CLONE_ERROR
-
-
-class GITConfigError(ArgumentInvalid):
-    """GIT configuration error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.GIT_CONFIG_ERROR
-
-
-class BlobStorageDownloadError(ClientError):
-    """Azcopy blobstorage download error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.BLOBSTORAGE_DOWNLOAD_ERROR
-
-
-class InvalidHuggingfaceModelIDError(ArgumentInvalid):
-    """Invalid Huggingface model ID error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.INVALID_HUGGING_FACE_MODEL_ID
-
-
-class HuggingFaceErrorInFetchingModelInfo(ConnectionFailure):
-    """Error in fetching model info."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.ERROR_FETCHING_HUGGING_FACE_MODEL_INFO
-
-
-class NonMsiAttachedComputeError(ArgumentInvalid):
-    """Internal Import Model Generic Error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.NON_MSI_ATTACHED_COMPUTE_ERROR
-
-
-class UserIdentityMissingError(ArgumentInvalid):
-    """Internal Import Model Generic Error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.USER_IDENTITY_MISSING_ERROR
-
-
-class VMNotSufficientForOperation(ArgumentInvalid):
-    """Error when VM is not sufficient for an operation."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.VM_NOT_SUFFICIENT_FOR_OPERATION
-
-
-class GenericRunCMDError(ClientError):
-    """Generic run CMD error."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.CMD_EXECUTION_ERROR
-
-
-class ModelAlreadyExists(ArgumentInvalid):
-    """Error when Model already exists in registry."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.MODEL_ALREADY_EXISTS
-
-
-class UnsupportedTaskType(NotSupported):
-    """Error when Unsupported task type is provided."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.UNSUPPORTED_TASK_TYPE
-
-
-class HFAuthenticationError(Authentication):
-    """Error when failed to authenticate user with token provided."""
-
-    @property
-    def message_format(self) -> str:
-        """Message format."""
-        return ModelImportErrorStrings.HF_AUTHENTICATION_ERROR
-
-
 def swallow_all_exceptions(logger: logging.Logger):
     """Swallow all exceptions.
 
@@ -208,10 +66,15 @@ def wrapper(*args, **kwargs):
             try:
                 return func(*args, **kwargs)
             except Exception as e:
-                if isinstance(e, AzureMLException):
+                if isinstance(e, MlException):
                     azureml_exception = e
                 else:
-                    azureml_exception = AzureMLException._with_error(AzureMLError.create(ModelImportError, error=e))
+                    message = ModelImportErrorStrings.LOG_UNSAFE_GENERIC_ERROR
+                    azureml_exception = MlException(
+                        message=message.format(error=e), no_personal_data_message=message,
+                        error_category=ErrorCategory.SYSTEM_ERROR, target=ErrorTarget.COMPONENT,
+                        error=e
+                    )
 
                 logger.error("Exception {} when calling {}".format(azureml_exception, func.__name__))
                 for handler in logger.handlers:

assets/training/model_management/src/run_model_download.py

@@ -8,11 +8,10 @@
 import re
 from azureml.model.mgmt.config import AppName, LlamaHFModels, LlamaModels, llama_dict
 from azureml.model.mgmt.downloader import download_model, ModelSource
-from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, ModelAlreadyExists
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, ModelImportErrorStrings
 from azureml.model.mgmt.utils.logging_utils import custom_dimensions, get_logger
 from azureml.model.mgmt.utils.common_utils import get_mlclient
-from azureml._common.exceptions import AzureMLException
-from azureml._common._error_definition.azureml_error import AzureMLError
+from azure.ai.ml.exceptions import ErrorTarget, ErrorCategory, MlException
 
 VALID_MODEL_NAME_PATTERN = r"^[a-zA-Z0-9-]+$"
 NEGATIVE_MODEL_NAME_PATTERN = r"[^a-zA-Z0-9-]"
@@ -72,8 +71,12 @@ def validate_if_model_exists(model_id):
         if model:
             version = model.version
             url = f"https://ml.azure.com/registries/{registry}/models/{model_id}/version/{version}"
-            raise AzureMLException._with_error(
-                AzureMLError.create(ModelAlreadyExists, model_id=model_id, registry=registry, url=url)
+            message = ModelImportErrorStrings.MODEL_ALREADY_EXISTS.format(
+                model_id=model_id, registry=registry, url=url
+            )
+            raise MlException(
+                message=message, no_personal_data_message=message,
+                error_category=ErrorCategory.USER_ERROR, target=ErrorTarget.COMPONENT
             )
         else:
             logger.info(f"Model {model_id} has not been imported into the registry. "

assets/training/model_management/src/run_model_preprocess.py

@@ -12,9 +12,8 @@
 from azureml.model.mgmt.processors.preprocess import run_preprocess, check_for_py_files
 from azureml.model.mgmt.processors.transformers.config import SupportedTasks as TransformersSupportedTasks
 from azureml.model.mgmt.processors.pyfunc.config import SupportedTasks as PyFuncSupportedTasks
-from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, UnsupportedTaskType
-from azureml._common.exceptions import AzureMLException
-from azureml._common._error_definition.azureml_error import AzureMLError
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, ModelImportErrorStrings
+from azure.ai.ml.exceptions import ErrorTarget, ErrorCategory, MlException
 from azureml.model.mgmt.utils.logging_utils import custom_dimensions, get_logger
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -134,10 +133,14 @@ def run():
 
     if task_name is None:
         supported_tasks = set(TransformersSupportedTasks.list_values() + PyFuncSupportedTasks.list_values())
-        raise AzureMLException._with_error(
-                AzureMLError.create(UnsupportedTaskType, task_type=args.task_name,
-                                    supported_tasks=list(supported_tasks))
-            )
+        message = ModelImportErrorStrings.UNSUPPORTED_TASK_TYPE.format(
+            task_type=args.task_name, supported_tasks=list(supported_tasks)
+        )
+        raise MlException(
+            message=message, no_personal_data_message=message,
+            error_category=ErrorCategory.USER_ERROR, target=ErrorTarget.COMPONENT
+        )
+
     files = check_for_py_files(model_path)
     logger.info(f"check if model folder contains .py files or not: {files}")
     if files: