Env dedup (Azure#38590)

Author: achauhan-scc

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_asset_utils.py

@@ -755,6 +755,39 @@ def _get_next_version_from_container(
     return version
 
 
+def _get_next_latest_versions_from_container(
+    name: str,
+    container_operation: Any,
+    resource_group_name: str,
+    workspace_name: str,
+    registry_name: str = None,
+    **kwargs,
+) -> str:
+    try:
+        container = (
+            container_operation.get(
+                name=name,
+                resource_group_name=resource_group_name,
+                registry_name=registry_name,
+                **kwargs,
+            )
+            if registry_name
+            else container_operation.get(
+                name=name,
+                resource_group_name=resource_group_name,
+                workspace_name=workspace_name,
+                **kwargs,
+            )
+        )
+        next_version = container.properties.next_version
+        latest_version = container.properties.latest_version
+
+    except ResourceNotFoundError:
+        next_version = "1"
+        latest_version = "1"
+    return next_version, latest_version
+
+
 def _get_latest_version_from_container(
     asset_name: str,
     container_operation: Any,

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py

@@ -26,7 +26,7 @@
 from azure.ai.ml._utils._asset_utils import (
     _archive_or_restore,
     _get_latest,
-    _get_next_version_from_container,
+    _get_next_latest_versions_from_container,
     _resolve_label_to_asset,
 )
 from azure.ai.ml._utils._experimental import experimental
@@ -39,7 +39,7 @@
 from azure.ai.ml.constants._common import ARM_ID_PREFIX, ASSET_ID_FORMAT, AzureMLResourceType
 from azure.ai.ml.entities._assets import Environment, WorkspaceAssetReference
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
-from azure.core.exceptions import ResourceNotFoundError
+from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
 
 ops_logger = OpsLogger(__name__)
 module_logger = ops_logger.module_logger
@@ -110,14 +110,18 @@ def create_or_update(self, environment: Environment) -> Environment:  # type: ig
         """
         try:
             if not environment.version and environment._auto_increment_version:
-                environment.version = _get_next_version_from_container(
+
+                next_version, latest_version = _get_next_latest_versions_from_container(
                     name=environment.name,
                     container_operation=self._containers_operations,
                     resource_group_name=self._operation_scope.resource_group_name,
                     workspace_name=self._workspace_name,
                     registry_name=self._registry_name,
                     **self._kwargs,
                 )
+                # If user not passing the version, SDK will try to update the latest version
+                return self._try_update_latest_version(next_version, latest_version, environment)
+
             sas_uri = None
             if self._registry_name:
                 if isinstance(environment, WorkspaceAssetReference):
@@ -204,6 +208,25 @@ def create_or_update(self, environment: Environment) -> Environment:  # type: ig
             else:
                 raise ex
 
+    def _try_update_latest_version(
+        self, next_version: str, latest_version: str, environment: Environment
+    ) -> Environment:
+        env = None
+        if self._registry_name:
+            environment.version = next_version
+            env = self.create_or_update(environment)
+        else:
+            environment.version = latest_version
+            try:  # Try to update the latest version
+                env = self.create_or_update(environment)
+            except Exception as ex:  # pylint: disable=W0718
+                if isinstance(ex, ResourceExistsError):
+                    environment.version = next_version
+                    env = self.create_or_update(environment)
+                else:
+                    raise ex
+        return env
+
     def _get(self, name: str, version: Optional[str] = None) -> EnvironmentVersion:
         if version:
             return (

sdk/ml/azure-ai-ml/tests/environment/unittests/test_environment_operations.py

@@ -1,5 +1,5 @@
 from typing import Iterable
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
 
 import pytest
 
@@ -14,7 +14,7 @@
 from azure.ai.ml.constants._common import ARM_ID_PREFIX
 from azure.ai.ml.entities._assets import Environment
 from azure.ai.ml.operations import EnvironmentOperations
-from azure.core.exceptions import ResourceNotFoundError
+from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
 
 
 @pytest.fixture
@@ -108,15 +108,106 @@ def test_create_autoincrement(
         with patch(
             "azure.ai.ml.operations._environment_operations.Environment._from_rest_object", return_value=None
         ), patch(
-            "azure.ai.ml.operations._environment_operations._get_next_version_from_container", return_value="version"
+            "azure.ai.ml.operations._environment_operations._get_next_latest_versions_from_container",
+            return_value=("version", "latest"),
         ) as mock_nextver:
             mock_environment_operation.create_or_update(env)
             mock_nextver.assert_called_once()
 
             mock_environment_operation._version_operations.create_or_update.assert_called_once_with(
                 body=env._to_rest_object(),
                 name=env.name,
-                version=mock_nextver.return_value,
+                version="latest",
+                resource_group_name=mock_workspace_scope.resource_group_name,
+                workspace_name=mock_workspace_scope.workspace_name,
+            )
+
+    def test_create_when_latest_fail_with_resource_exist(
+        self,
+        mock_environment_operation: EnvironmentOperations,
+        mock_workspace_scope: OperationScope,
+    ) -> None:
+        env = load_environment(source="./tests/test_configs/environment/environment_no_version.yml")
+        assert env._auto_increment_version
+        env.version = None
+        mock_environment_operation._version_operations.create_or_update(
+            body=env._to_rest_object(),
+            name=env.name,
+            version="latest_version",
+            resource_group_name=mock_workspace_scope.resource_group_name,
+            workspace_name=mock_workspace_scope.workspace_name,
+        ).side_effect = ResourceExistsError()
+        with patch(
+            "azure.ai.ml.operations._environment_operations.Environment._from_rest_object", return_value=None
+        ), patch(
+            "azure.ai.ml.operations._environment_operations._get_next_latest_versions_from_container",
+            return_value=("next_version", "latest_version"),
+        ) as mock_nextver:
+
+            def side_effect(name, version, body, resource_group_name, workspace_name):
+                if version == "latest_version":
+                    raise ResourceExistsError
+
+            mock_environment_operation._version_operations.create_or_update.side_effect = side_effect
+            mock_environment_operation.create_or_update(env)
+            mock_nextver.assert_called_once()
+
+            mock_environment_operation._version_operations.create_or_update.assert_has_calls(
+                [
+                    call(
+                        body=env._to_rest_object(),
+                        name=env.name,
+                        version="latest_version",
+                        resource_group_name=mock_workspace_scope.resource_group_name,
+                        workspace_name=mock_workspace_scope.workspace_name,
+                    ),
+                    call(
+                        body=env._to_rest_object(),
+                        name=env.name,
+                        version="next_version",
+                        resource_group_name=mock_workspace_scope.resource_group_name,
+                        workspace_name=mock_workspace_scope.workspace_name,
+                    ),
+                ]
+            )
+
+    def test_create_when_latest_fail_with_unknown(
+        self,
+        mock_environment_operation: EnvironmentOperations,
+        mock_workspace_scope: OperationScope,
+    ) -> None:
+        env = load_environment(source="./tests/test_configs/environment/environment_no_version.yml")
+        assert env._auto_increment_version
+        env.version = None
+        mock_environment_operation._version_operations.create_or_update(
+            body=env._to_rest_object(),
+            name=env.name,
+            version="latest_version",
+            resource_group_name=mock_workspace_scope.resource_group_name,
+            workspace_name=mock_workspace_scope.workspace_name,
+        ).side_effect = ResourceExistsError()
+        with patch(
+            "azure.ai.ml.operations._environment_operations.Environment._from_rest_object", return_value=None
+        ), patch(
+            "azure.ai.ml.operations._environment_operations._get_next_latest_versions_from_container",
+            return_value=("next_version", "latest_version"),
+        ) as mock_nextver:
+
+            def side_effect(name, version, body, resource_group_name, workspace_name):
+                if version == "latest_version":
+                    raise Exception("Unknowm Error")
+
+            mock_environment_operation._version_operations.create_or_update.side_effect = side_effect
+            try:
+                mock_environment_operation.create_or_update(env)
+            except Exception:
+                pass
+            mock_nextver.assert_called_once()
+
+            mock_environment_operation._version_operations.create_or_update.assert_called_with(
+                body=env._to_rest_object(),
+                name=env.name,
+                version="latest_version",
                 resource_group_name=mock_workspace_scope.resource_group_name,
                 workspace_name=mock_workspace_scope.workspace_name,
             )
@@ -177,6 +268,9 @@ def test_restore_container(self, mock_environment_operation: EnvironmentOperatio
             resource_group_name=mock_environment_operation._resource_group_name,
         )
 
+    def side_effect(self, args):
+        print(args)
+
     # #Mock(azure.ai.ml._restclient.v2021_10_01_dataplanepreview.operations._environment_versions_operations, "get")
     # def test_promote_environment_from_workspace(
     #     self,