Add registry replication count value (Azure#26932)

* schema changes an initial support class changes

* context switch changes

* swagger changes

* update tests and add tests for replicated ids

* nits

* fix lint issues

* linting pass

* pylint second pass

* fix merge

* get e2e test passing again after merge

* more testing

* remove pdb

* linting

* comment nits

* rerun test recordings after main merge

* remove description

* schema deserialization fix

* fix lint

* add replication count check at conversion

Author: MilesHolland

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/registry/system_created_storage_account.py

@@ -15,10 +15,13 @@
 @experimental
 class SystemCreatedStorageAccountSchema(metaclass=PatchedSchemaMeta):
     arm_resource_id = fields.Str(dump_only=True)
-    storage_account_hns = fields.Bool()
+    storage_account_hns = fields.Bool(load_default=False)
     storage_account_type = StringTransformedEnum(
+        load_default=StorageAccountType.STANDARD_LRS,
         allowed_values=[accountType.value for accountType in StorageAccountType], casing_transform=lambda x: x.lower()
     )
+    replication_count = fields.Int(load_default=1, validate=lambda count : count > 0)
+    replicated_ids = fields.List(fields.Str(), dump_only=True)
 
     @post_load
     def make(self, data, **kwargs):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_registry/registry.py

@@ -21,7 +21,7 @@
 from azure.ai.ml.entities._util import load_from_dict
 from azure.ai.ml._utils._experimental import experimental
 
-from .registry_support_classes import RegistryRegionDetails, SystemCreatedStorageAccount
+from .registry_support_classes import RegistryRegionDetails
 
 CONTAINER_REGISTRY = "container_registry"
 REPLICATION_LOCATIONS = "replication_locations"
@@ -34,7 +34,6 @@ def __init__(
         name: str,
         location: str,
         identity: IdentityConfiguration = None,
-        description: str = None,
         tags: Dict[str, str] = None,
         public_network_access: str = None,
         discovery_url: str = None,
@@ -52,8 +51,6 @@ def __init__(
         :type location: str
         :param identity: registry's System Managed Identity
         :type identity: ManagedServiceIdentity
-        :param description: Description of the registry.
-        :type description: str
         :param tags: Tags of the registry.
         :type tags: dict
         :param public_network_access: Whether to allow public endpoint connectivity.
@@ -72,7 +69,7 @@ def __init__(
         :type kwargs: dict
         """
 
-        super().__init__(name=name, description=description, tags=tags, **kwargs)
+        super().__init__(name=name, tags=tags, **kwargs)
 
         # self.display_name = name # Do we need a top-level visible name value?
         self.location = location
@@ -118,15 +115,6 @@ def _to_dict(self) -> Dict:
             if self.replication_locations[0].acr_config and len(self.replication_locations[0].acr_config) > 0:
                 self.container_registry = self.replication_locations[0].acr_config[0]
 
-        # Change single-list managed storage accounts to not be lists.
-        # Although storage accounts are storage in a list to match the
-        # underlying API, users should only enter in one managed storage
-        # in YAML.
-        for region_detail in self.replication_locations:
-            if region_detail.storage_config and isinstance(
-                region_detail.storage_config[0], SystemCreatedStorageAccount
-            ):
-                region_detail.storage_config = region_detail.storage_config[0]
         return schema.dump(self)
 
     @classmethod
@@ -166,7 +154,6 @@ def _from_rest_object(cls, rest_obj: RestRegistry) -> "Registry":
             identity = IdentityConfiguration._from_rest_object(rest_obj.identity)
         return Registry(
             name=rest_obj.name,
-            description=real_registry.description,
             identity=identity,
             id=rest_obj.id,
             tags=rest_obj.tags,
@@ -200,10 +187,6 @@ def _convert_yaml_dict_to_entity_input(
             if global_acr_exists:
                 if not hasattr(region_detail, "acr_details") or len(region_detail.acr_details) == 0:
                     region_detail.acr_config = [acr_input]
-            # Convert single, non-list managed storage into a 1-element list.
-            if hasattr(region_detail, "storage_config") and isinstance(region_detail.storage_config, \
-                                                                        SystemCreatedStorageAccount):
-                region_detail.storage_config = [region_detail.storage_config]
 
 
     def _to_rest_object(self) -> RestRegistry:
@@ -227,7 +210,6 @@ def _to_rest_object(self) -> RestRegistry:
             location=self.location,
             identity=identity,
             tags=self.tags,
-            description=self.description,
             properties=RegistryProperties(
                 public_network_access=self.public_network_access,
                 discovery_url=self.discovery_url,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_registry/registry_support_classes.py

@@ -5,18 +5,21 @@
 
 from typing import List, Union
 
+from functools import reduce
+from copy import deepcopy
 from azure.ai.ml._restclient.v2022_10_01_preview.models import (
     AcrDetails as RestAcrDetails,
     ArmResourceId as RestArmResourceId,
     RegistryRegionArmDetails as RestRegistryRegionArmDetails,
     StorageAccountDetails as RestStorageAccountDetails,
     SystemCreatedAcrAccount as RestSystemCreatedAcrAccount,
     SystemCreatedStorageAccount as RestSystemCreatedStorageAccount,
-    UserCreatedAcrAccount as RestUserCreatedAcrAccount,
-    UserCreatedStorageAccount as RestUserCreatedStorageAccount)
+    UserCreatedAcrAccount as RestUserCreatedAcrAccount)
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._registry import StorageAccountType
-
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
+from azure.ai.ml._exception_helper import log_and_raise_error
+from .util import _make_rest_user_storage_from_id
 
 # This exists despite not being used by the schema validator because this entire
 # class is an output only value from the API.
@@ -93,6 +96,8 @@ def __init__(
         storage_account_hns: bool,
         storage_account_type: StorageAccountType,
         arm_resource_id: str = None,
+        replicated_ids: List[str] = None,
+        replication_count = 1,
     ):
         """
         :param arm_resource_id: Resource ID of the storage account.
@@ -104,58 +109,19 @@ def __init__(
             "Standard_GRS, "Standard_RAGRS", "Standard_ZRS", "Standard_GZRS",
             "Standard_RAGZRS", "Premium_LRS", "Premium_ZRS"
         :type storage_account_type: StorageAccountType
+        :param replication_count: The number of replicas of this storage account
+            that should be created. Defaults to 1. Values less than 1 are invalid.
+        :type replication_count: int
+        :param replicated_ids: If this storage was replicated, then this is a
+            list of all storage IDs with these settings for this registry.
+            Defaults to none for un-replicated storage accounts.
+        :type replicated_ids: List[str]
         """
         self.arm_resource_id = arm_resource_id
         self.storage_account_hns = storage_account_hns
         self.storage_account_type = storage_account_type
-
-    # storage should technically be a union between str and SystemCreatedStorageAccount,
-    # but python doesn't accept self class references apparently.
-    # Class method instead of normal function to accept possible
-    # string input.
-
-    @classmethod
-    def _to_rest_object(cls, storage) -> RestStorageAccountDetails:
-        if hasattr(storage, "storage_account_type") and storage.storage_account_type is not None:
-
-            # We DO NOT want to set the arm_resource_id. The backend provides very
-            # unhelpful errors if you provide an empty/null/invalid resource ID,
-            # and ignores the value otherwise. It's better to avoid setting it in
-            # the conversion in this direction at all.
-            # We don't bother processing storage_account_type because the
-            # rest version is case insensitive.
-            account = RestSystemCreatedStorageAccount(
-                storage_account_hns_enabled=storage.storage_account_hns,
-                storage_account_type=storage.storage_account_type,
-            )
-            return RestStorageAccountDetails(system_created_storage_account=account)
-        else:
-            return RestStorageAccountDetails(
-                user_created_storage_account=RestUserCreatedStorageAccount(
-                    arm_resource_id=RestArmResourceId(resource_id=storage)
-                )
-            )
-
-    @classmethod
-    def _from_rest_object(cls, rest_obj: RestStorageAccountDetails) -> "Union[str, SystemCreatedStorageAccount]":
-        if not rest_obj:
-            return None
-        # TODO should we even bother check if both values are set and throw an error? This shouldn't be possible.
-        if rest_obj.system_created_storage_account:
-            resource_id = None
-            if rest_obj.system_created_storage_account.arm_resource_id:
-                resource_id = rest_obj.system_created_storage_account.arm_resource_id.resource_id
-            return SystemCreatedStorageAccount(
-                storage_account_hns=rest_obj.system_created_storage_account.storage_account_hns_enabled,
-                storage_account_type=StorageAccountType(
-                    rest_obj.system_created_storage_account.storage_account_type.lower()
-                ),  # TODO validate storage account type?
-                arm_resource_id=resource_id,
-            )
-        elif rest_obj.user_created_storage_account:
-            return rest_obj.user_created_storage_account.arm_resource_id.resource_id
-        else:
-            return None  # TODO should this throw an error instead?
+        self.replication_count = replication_count
+        self.replicated_ids = replicated_ids
 
 
 # Per-region information for registries.
@@ -166,7 +132,7 @@ def __init__(
         *,
         acr_config: List[Union[str, SystemCreatedAcrAccount]] = None,
         location: str = None,
-        storage_config: List[Union[str, SystemCreatedStorageAccount]] = None,
+        storage_config: Union[List[str], SystemCreatedStorageAccount] = None,
     ):
         """
         Details for each region a registry is in.
@@ -197,10 +163,12 @@ def _from_rest_object(cls, rest_obj: RestRegistryRegionArmDetails) -> "RegistryR
                 acr) for acr in rest_obj.acr_details]
         storages = []
         if rest_obj.storage_account_details:
-            storages = [SystemCreatedStorageAccount._from_rest_object(
-                storages) for storages in rest_obj.storage_account_details]
+            storages = cls._storage_config_from_rest_object(rest_obj.storage_account_details)
+
         return RegistryRegionDetails(
-            acr_config=converted_acr_details, location=rest_obj.location, storage_config=storages
+            acr_config=converted_acr_details,
+            location=rest_obj.location,
+            storage_config=storages
         )
 
     def _to_rest_object(self) -> RestRegistryRegionArmDetails:
@@ -210,10 +178,91 @@ def _to_rest_object(self) -> RestRegistryRegionArmDetails:
                 acr) for acr in self.acr_config]
         storages = []
         if self.storage_config:
-            storages = [SystemCreatedStorageAccount._to_rest_object(
-                storage) for storage in self.storage_config]
+            storages = self._storage_config_to_rest_object()
         return RestRegistryRegionArmDetails(
             acr_details=converted_acr_details,
             location=self.location,
             storage_account_details=storages,
         )
+
+
+    def _storage_config_to_rest_object(self) -> List[RestStorageAccountDetails]:
+        storage = self.storage_config
+        # storage_config can either be a single system-created storage account,
+        # or list of user-inputted id's.
+        if (storage is not None
+            and hasattr(storage, "storage_account_type")
+            and storage.storage_account_type is not None):
+
+            # We DO NOT want to set the arm_resource_id. The backend provides very
+            # unhelpful errors if you provide an empty/null/invalid resource ID,
+            # and ignores the value otherwise. It's better to avoid setting it in
+            # the conversion in this direction at all.
+            # We don't bother processing storage_account_type because the
+            # rest version is case insensitive.
+            account = RestStorageAccountDetails(system_created_storage_account=RestSystemCreatedStorageAccount(
+                storage_account_hns_enabled=storage.storage_account_hns,
+                storage_account_type=storage.storage_account_type,
+            ))
+            # duplicate this value based on the replication_count
+            count = storage.replication_count
+            if count < 1:
+                raise ValueError(f"Replication count cannot be less than 1. Value was: {count}.")
+            return [deepcopy(account) for _ in range(0, count)]
+        elif storage is not None and len(storage) > 0:
+            return [_make_rest_user_storage_from_id(user_id=user_id) for user_id in storage]
+        else:
+            return []
+
+    @classmethod
+    def _storage_config_from_rest_object(
+        cls,
+        rest_configs: List[RestStorageAccountDetails]
+    ) -> Union[List[str], SystemCreatedStorageAccount]:
+        if not rest_configs:
+            return None
+        num_configs = len(rest_configs)
+        if num_configs == 0:
+            return None
+        system_created_count = reduce(lambda x, y: int(x) + int(y),
+            [hasattr(config, "system_created_storage_account") and config.system_created_storage_account is not None
+            for config in rest_configs])
+        # configs should be mono-typed. Either they're all system created
+        # or all user created.
+        if system_created_count == num_configs:
+            # System created case - assume all elements are duplicates
+            # of a single storage configuration.
+            # Convert back into a single local representation by
+            # combining id's into a list, and using the first element's
+            # account type and hns.
+            first_config = rest_configs[0].system_created_storage_account
+            resource_id = None
+            if first_config.arm_resource_id:
+                resource_id = first_config.arm_resource_id.resource_id
+            # account for ids of duplicated if they exist
+            replicated_ids = None
+            if num_configs > 1:
+                replicated_ids = [config.system_created_storage_account.arm_resource_id.resource_id
+                    for config in rest_configs]
+            return SystemCreatedStorageAccount(
+                storage_account_hns=first_config.storage_account_hns_enabled,
+                storage_account_type=StorageAccountType(
+                    first_config.storage_account_type.lower()
+                ),  # TODO validate storage account type? GI
+                arm_resource_id=resource_id,
+                replication_count=num_configs,
+                replicated_ids=replicated_ids
+            )
+        elif  system_created_count == 0:
+            return [config.user_created_storage_account.arm_resource_id.resource_id for config in rest_configs]
+        else:
+            msg = f'''tried reading in a registry whose storage accounts were not
+                mono-managed or user-created. {system_created_count} out of {num_configs} were managed.'''
+            err = ValidationException(
+                message=msg,
+                target=ErrorTarget.REGISTRY,
+                no_personal_data_message=msg,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
+            log_and_raise_error(err)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_registry/util.py

@@ -0,0 +1,16 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+# pylint:disable=protected-access,no-else-return
+
+from azure.ai.ml._restclient.v2022_10_01_preview.models import (
+    ArmResourceId as RestArmResourceId,
+    StorageAccountDetails as RestStorageAccountDetails,
+    UserCreatedStorageAccount as RestUserCreatedStorageAccount)
+
+
+def _make_rest_user_storage_from_id(*, user_id: str) -> RestStorageAccountDetails:
+    return RestStorageAccountDetails(user_created_storage_account=RestUserCreatedStorageAccount(
+            arm_resource_id=RestArmResourceId(resource_id=user_id)
+        )
+    )