[ml] Support of Environment promotion from workspace to registry (Azure#28588)

banibrata-de · web-flow · commit 5baed21fca5b · 2023-02-01T18:45:42.000-08:00
* promo of environment tests
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py
@@ -57,7 +57,7 @@
     Registry,
     Workspace,
 )
-from azure.ai.ml.entities._assets import WorkspaceModelReference
+from azure.ai.ml.entities._assets import WorkspaceAssetReference
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
 from azure.ai.ml.operations import (
     BatchDeploymentOperations,
@@ -758,7 +758,7 @@ def _get_workspace_info(cls, found_path: Optional[str]) -> Tuple[str, str, str]:
 
     # T = valid inputs/outputs for create_or_update
     # Each entry here requires a registered _create_or_update function below
-    T = TypeVar("T", Job, Model, Environment, Component, Datastore, WorkspaceModelReference)
+    T = TypeVar("T", Job, Model, Environment, Component, Datastore, WorkspaceAssetReference)
 
     def create_or_update(
         self,
@@ -770,7 +770,7 @@ def create_or_update(
         :param entity: The resource to create or update.
         :type entity: typing.Union[~azure.ai.ml.entities.Job
             , ~azure.ai.ml.entities.Model, ~azure.ai.ml.entities.Environment, ~azure.ai.ml.entities.Component
-            , ~azure.ai.ml.entities.Datastore, ~azure.ai.ml.entities.WorkspaceModelReference]
+            , ~azure.ai.ml.entities.Datastore, ~azure.ai.ml.entities.WorkspaceAssetReference]
         :return: The created or updated resource.
         :rtype: typing.Union[~azure.ai.ml.entities.Job, ~azure.ai.ml.entities.Model
             , ~azure.ai.ml.entities.Environment, ~azure.ai.ml.entities.Component, ~azure.ai.ml.entities.Datastore]
@@ -835,8 +835,8 @@ def _(entity: Model, operations):
     return operations[AzureMLResourceType.MODEL].create_or_update(entity)
 
 
-@_create_or_update.register(WorkspaceModelReference)
-def _(entity: WorkspaceModelReference, operations):
+@_create_or_update.register(WorkspaceAssetReference)
+def _(entity: WorkspaceAssetReference, operations):
     module_logger.debug("Promoting model to registry")
     return operations[AzureMLResourceType.MODEL].create_or_update(entity)
 
@@ -847,6 +847,12 @@ def _(entity: Environment, operations):
     return operations[AzureMLResourceType.ENVIRONMENT].create_or_update(entity)
 
 
+@_create_or_update.register(WorkspaceAssetReference)
+def _(entity: WorkspaceAssetReference, operations):
+    module_logger.debug("Promoting environment to registry")
+    return operations[AzureMLResourceType.ENVIRONMENT].create_or_update(entity)
+
+
 @_create_or_update.register(Component)
 def _(entity: Component, operations, **kwargs):
     module_logger.debug("Creating or updating components")
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/__init__.py
@@ -9,7 +9,7 @@
 from .assets.data import DataSchema
 from .assets.environment import AnonymousEnvironmentSchema, EnvironmentSchema
 from .assets.model import ModelSchema
-from .assets.workspace_model_reference import WorkspaceModelReferenceSchema
+from .assets.workspace_asset_reference import WorkspaceAssetReferenceSchema
 from .component import CommandComponentSchema
 from .core.fields import (
     ArmStr,
@@ -52,5 +52,5 @@
     "AnonymousCodeAssetSchema",
     "ExperimentalField",
     "RegistryStr",
-    "WorkspaceModelReferenceSchema",
+    "WorkspaceAssetReferenceSchema",
 ]
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/assets/workspace_asset_reference.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/assets/workspace_asset_reference.py
@@ -15,13 +15,13 @@
 module_logger = logging.getLogger(__name__)
 
 
-class WorkspaceModelReferenceSchema(AssetSchema):
+class WorkspaceAssetReferenceSchema(AssetSchema):
     destination_name = fields.Str()
     destination_version = fields.Str()
     source_asset_id = fields.Str(required=True)
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._assets import WorkspaceModelReference
+        from azure.ai.ml.entities._assets import WorkspaceAssetReference
 
-        return WorkspaceModelReference(base_path=self.context[BASE_PATH_CONTEXT_KEY], **data)
+        return WorkspaceAssetReference(base_path=self.context[BASE_PATH_CONTEXT_KEY], **data)
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py
@@ -15,7 +15,7 @@
 from ._assets._artifacts.model import Model
 from ._assets.asset import Asset
 from ._assets.environment import BuildContext, Environment
-from ._assets.workspace_model_reference import WorkspaceModelReference
+from ._assets.workspace_asset_reference import WorkspaceAssetReference
 from ._builders import Command, Parallel, Pipeline, Spark, Sweep
 from ._component.command_component import CommandComponent
 from ._component.component import Component
@@ -246,7 +246,7 @@
     "SynapseSparkCompute",
     "AutoScaleSettings",
     "AutoPauseSettings",
-    "WorkspaceModelReference",
+    "WorkspaceAssetReference",
     # builders
     "Command",
     "Parallel",
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/__init__.py
@@ -10,6 +10,6 @@
 from ._artifacts.data import Data
 from ._artifacts.model import Model
 from .environment import Environment
-from .workspace_model_reference import WorkspaceModelReference
+from .workspace_asset_reference import WorkspaceAssetReference
 
-__all__ = ["Artifact", "Model", "Code", "Data", "Environment", "WorkspaceModelReference"]
+__all__ = ["Artifact", "Model", "Code", "Data", "Environment", "WorkspaceAssetReference"]
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/workspace_asset_reference.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/workspace_asset_reference.py
@@ -10,13 +10,13 @@
     ResourceManagementAssetReferenceData,
     ResourceManagementAssetReferenceDetails,
 )
-from azure.ai.ml._schema import WorkspaceModelReferenceSchema
+from azure.ai.ml._schema import WorkspaceAssetReferenceSchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY
 from azure.ai.ml.entities._assets.asset import Asset
 from azure.ai.ml.entities._util import load_from_dict
 
 
-class WorkspaceModelReference(Asset):
+class WorkspaceAssetReference(Asset):
     """Workspace Model Reference.
 
     :param name: Model name
@@ -53,14 +53,14 @@ def _load(
         yaml_path: Optional[Union[os.PathLike, str]] = None,
         params_override: Optional[list] = None,
         **kwargs,
-    ) -> "WorkspaceModelReference":
+    ) -> "WorkspaceAssetReference":
         data = data or {}
         params_override = params_override or []
         context = {
             BASE_PATH_CONTEXT_KEY: Path(yaml_path).parent if yaml_path else Path("./"),
             PARAMS_OVERRIDE_KEY: params_override,
         }
-        return load_from_dict(WorkspaceModelReferenceSchema, data, context, **kwargs)
+        return load_from_dict(WorkspaceAssetReferenceSchema, data, context, **kwargs)
 
     def _to_rest_object(self) -> ResourceManagementAssetReferenceData:
         resource_management_details = ResourceManagementAssetReferenceDetails(
@@ -72,9 +72,9 @@ def _to_rest_object(self) -> ResourceManagementAssetReferenceData:
         return resource_management
 
     @classmethod
-    def _from_rest_object(cls, resource_object: ResourceManagementAssetReferenceData) -> "WorkspaceModelReference":
+    def _from_rest_object(cls, resource_object: ResourceManagementAssetReferenceData) -> "WorkspaceAssetReference":
 
-        resource_management = WorkspaceModelReference(
+        resource_management = WorkspaceAssetReference(
             name=resource_object.properties.destination_name,
             version=resource_object.properties.destination_version,
             asset_id=resource_object.properties.source_asset_id,
@@ -84,4 +84,4 @@ def _from_rest_object(cls, resource_object: ResourceManagementAssetReferenceData
 
     def _to_dict(self) -> Dict:
         # pylint: disable=no-member
-        return WorkspaceModelReferenceSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self)
+        return WorkspaceAssetReferenceSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self)
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py
@@ -31,8 +31,9 @@
 )
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml._utils._registry_utils import get_asset_body_for_registry_storage, get_sas_uri_for_registry_asset
-from azure.ai.ml.constants._common import ARM_ID_PREFIX, AzureMLResourceType
-from azure.ai.ml.entities._assets import Environment
+from azure.ai.ml.constants._common import ARM_ID_PREFIX, AzureMLResourceType, ASSET_ID_FORMAT
+from azure.ai.ml.entities._assets import Environment, WorkspaceAssetReference
+from azure.core.exceptions import ResourceNotFoundError
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
 ops_logger = OpsLogger(__name__)
@@ -92,6 +93,37 @@ def create_or_update(self, environment: Environment) -> Environment:
                 )
             sas_uri = None
             if self._registry_name:
+                if isinstance(environment, WorkspaceAssetReference):
+                    # verify that environment is not already in registry
+                    try:
+                        self._version_operations.get(
+                            name=environment.name,
+                            version=environment.version,
+                            resource_group_name=self._resource_group_name,
+                            registry_name=self._registry_name,
+                        )
+                    except Exception as err:  # pylint: disable=broad-except
+                        if isinstance(err, ResourceNotFoundError):
+                            pass
+                        else:
+                            raise err
+                    else:
+                        msg = "A environment with this name and version already exists in registry"
+                        raise ValidationException(
+                            message=msg,
+                            no_personal_data_message=msg,
+                            target=ErrorTarget.ENVIRONMENT,
+                            error_category=ErrorCategory.USER_ERROR,
+                        )
+
+                    environment = environment._to_rest_object()
+                    result = self._service_client.resource_management_asset_reference.begin_import_method(
+                        resource_group_name=self._resource_group_name,
+                        registry_name=self._registry_name,
+                        body=environment,
+                    )
+                    return result
+
                 sas_uri = get_sas_uri_for_registry_asset(
                     service_client=self._service_client,
                     name=environment.name,
@@ -111,9 +143,7 @@ def create_or_update(self, environment: Environment) -> Environment:
                     )
                     return self.get(name=environment.name, version=environment.version)
 
-            environment = _check_and_upload_env_build_context(
-                environment=environment, operations=self, sas_uri=sas_uri
-            )
+            environment = _check_and_upload_env_build_context(environment=environment, operations=self, sas_uri=sas_uri)
             env_version_resource = environment._to_rest_object()
             env_rest_obj = (
                 self._version_operations.begin_create_or_update(
@@ -279,7 +309,7 @@ def archive(
         name: str,
         version: Optional[str] = None,
         label: Optional[str] = None,
-        **kwargs # pylint:disable=unused-argument
+        **kwargs,  # pylint:disable=unused-argument
     ) -> None:
         """Archive an environment or an environment version.
 
@@ -307,7 +337,7 @@ def restore(
         name: str,
         version: Optional[str] = None,
         label: Optional[str] = None,
-        **kwargs # pylint:disable=unused-argument
+        **kwargs,  # pylint:disable=unused-argument
     ) -> None:
         """Restore an archived environment version.
 
@@ -343,6 +373,43 @@ def _get_latest_version(self, name: str) -> Environment:
         )
         return Environment._from_rest_object(result)
 
+    # pylint: disable=no-self-use
+    def _prepare_to_copy(
+        self, environment: Environment, name: Optional[str] = None, version: Optional[str] = None
+    ) -> WorkspaceAssetReference:
+
+        """Returns WorkspaceAssetReference
+        to copy a registered environment to registry given the asset id
+
+        :param environment: Registered environment
+        :type environment: Environment
+        :param name: Destination name
+        :type name: str
+        :param version: Destination version
+        :type version: str
+        """
+        #  Get workspace info to get workspace GUID
+        workspace = self._service_client.workspaces.get(
+            resource_group_name=self._resource_group_name, workspace_name=self._workspace_name
+        )
+        workspace_guid = workspace.workspace_id
+        workspace_location = workspace.location
+
+        # Get environment asset ID
+        asset_id = ASSET_ID_FORMAT.format(
+            workspace_location,
+            workspace_guid,
+            AzureMLResourceType.ENVIRONMENT,
+            environment.name,
+            environment.version,
+        )
+
+        return WorkspaceAssetReference(
+            name=name if name else environment.name,
+            version=version if version else environment.version,
+            asset_id=asset_id,
+        )
+
 
 def _preprocess_environment_name(environment_name: str) -> str:
     if environment_name.startswith(ARM_ID_PREFIX):
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py
@@ -43,7 +43,7 @@
 from azure.ai.ml._utils._storage_utils import get_ds_name_and_path_prefix, get_storage_client
 from azure.ai.ml._utils.utils import resolve_short_datastore_url, validate_ml_flow_folder
 from azure.ai.ml.constants._common import ASSET_ID_FORMAT, AzureMLResourceType
-from azure.ai.ml.entities._assets import Model, WorkspaceModelReference
+from azure.ai.ml.entities._assets import Model, WorkspaceAssetReference
 from azure.ai.ml.entities._credentials import AccountKeyConfiguration
 from azure.ai.ml.exceptions import (
     AssetPathException,
@@ -89,7 +89,7 @@ def __init__(
 
     # @monitor_with_activity(logger, "Model.CreateOrUpdate", ActivityType.PUBLICAPI)
     def create_or_update(
-        self, model: Union[Model, WorkspaceModelReference]
+        self, model: Union[Model, WorkspaceAssetReference]
     ) -> Model:  # TODO: Are we going to implement job_name?
         """Returns created or updated model asset.
 
@@ -120,7 +120,7 @@ def create_or_update(
 
             if self._registry_name:
                 # Case of copy model to registry
-                if isinstance(model, WorkspaceModelReference):
+                if isinstance(model, WorkspaceAssetReference):
                     # verify that model is not already in registry
                     try:
                         self._model_versions_operation.get(
@@ -464,9 +464,9 @@ def _get_latest_version(self, name: str) -> Model:
     # pylint: disable=no-self-use
     def _prepare_to_copy(
         self, model: Model, name: Optional[str] = None, version: Optional[str] = None
-    ) -> WorkspaceModelReference:
+    ) -> WorkspaceAssetReference:
 
-        """Returns WorkspaceModelReference
+        """Returns WorkspaceAssetReference
         to copy a registered model to registry given the asset id
 
         :param model: Registered model
@@ -492,7 +492,7 @@ def _prepare_to_copy(
             model.version,
         )
 
-        return WorkspaceModelReference(
+        return WorkspaceAssetReference(
             name=name if name else model.name,
             version=version if version else model.version,
             asset_id=asset_id,
diff --git a/sdk/ml/azure-ai-ml/tests/environment/unittests/test_environment_operations.py b/sdk/ml/azure-ai-ml/tests/environment/unittests/test_environment_operations.py
