Refactoring of single VMSS code (#393)

* Refactoring of single VMSS code

Co-authored-by: Karthick Narendran <kanarend@microsoft.com>
diff --git a/ansible/library/azure_rm_virtualmachinescaleset_nic_list_facts.py b/ansible/library/azure_rm_virtualmachinescaleset_nic_list_facts.py
new file mode 100644
index 0000000..2244b5a
--- /dev/null
+++ b/ansible/library/azure_rm_virtualmachinescaleset_nic_list_facts.py
@@ -0,0 +1,289 @@
+#!/usr/bin/python3
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# flake8: noqa
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+import six
+
+
+ANSIBLE_METADATA = {'metadata_version': '1.1',
+                    'status': ['preview'],
+                    'supported_by': 'community'}
+
+
+DOCUMENTATION = '''
+---
+module: azure_rm_virtualmachinescaleset_nic_list_facts
+
+version_added: "2.9"
+
+short_description: Get network interface facts for a VMSS
+
+description:
+    - Get a list of all network interfaces for a virtual machine scale set
+
+options:
+    vmss_name:
+        description:
+            - Name of the virtual machine scale set.
+        required: true
+    resource_group:
+        description:
+            - Name of the resource group containing the virtual machine scale set.
+        required: true
+
+extends_documentation_fragment:
+    - azure
+
+author:
+    - "Min Pae (@sputnik13)"
+
+'''
+
+EXAMPLES = '''
+    - name: Get all network interfaces in a virtual machine scale set
+      azure_rm_virtualmachinescaleset_nic_list_facts:
+        resource_group: myResourceGroup
+        vmss_name: myvmss
+'''
+
+RETURN = '''
+azure_networkinterfaces:
+    description: List of network interface dicts.
+    returned: always
+    type: list
+    example: [{
+        "dns_settings": {
+            "applied_dns_servers": [],
+            "dns_servers": [],
+            "internal_dns_name_label": null,
+            "internal_fqdn": null
+        },
+        "enable_ip_forwarding": false,
+        "etag": 'W/"59726bfc-08c4-44ed-b900-f6a559876a9d"',
+        "id": "/subscriptions/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx/resourceGroup/myResourceGroup/providers/Microsoft.Network/networkInterfaces/nic003",
+        "ip_configuration": {
+            "name": "default",
+            "private_ip_address": "10.10.0.4",
+            "private_ip_allocation_method": "Dynamic",
+            "public_ip_address": {
+                "id": "/subscriptions/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx/resourceGroup/myResourceGroup/providers/Microsoft.Network/publicIPAddresses/publicip001",
+                "name": "publicip001"
+            },
+            "subnet": {
+                "id": "/subscriptions/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx/resourceGroup/myResourceGroup/providers/Microsoft.Network/virtualNetworks/vnet001/subnets/subnet001",
+                "name": "subnet001",
+                "virtual_network_name": "vnet001"
+            }
+        },
+        "location": "westus",
+        "mac_address": null,
+        "name": "nic003",
+        "network_security_group": {
+            "id": "/subscriptions/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx/resourceGroup/myResourceGroup/providers/Microsoft.Network/networkSecurityGroups/secgroup001",
+            "name": "secgroup001"
+        },
+        "primary": null,
+        "provisioning_state": "Succeeded",
+        "tags": {},
+        "type": "Microsoft.Network/networkInterfaces"
+    }]
+networkinterfaces:
+    description: List of network interface dict, the dict contains parameters can be passed to C(azure_rm_networkinterface) module.
+    type: list
+    returned: always
+    contains:
+        id:
+            description:
+                - Id of the network interface.
+        resource_group:
+            description:
+                - Name of a resource group where the network interface exists.
+        name:
+            description:
+                - Name of the network interface.
+        location:
+            description:
+                - Azure location.
+        virtual_network:
+            description:
+                - An existing virtual network with which the network interface will be associated.
+                - It is a dict which contains C(name) and C(resource_group) of the virtual network.
+        subnet:
+            description:
+                - Name of an existing subnet within the specified virtual network.
+        tags:
+            description:
+                - Tags of the network interface.
+        ip_configurations:
+            description:
+                - List of ip configuration if contains mutilple configuration.
+            contains:
+                name:
+                    description:
+                        - Name of the ip configuration.
+                private_ip_address:
+                    description:
+                        - Private ip address for the ip configuration.
+                private_ip_allocation_method:
+                    description:
+                        - private ip allocation method.
+                public_ip_address:
+                    description:
+                        - Name of the public ip address. None for disable ip address.
+                public_ip_allocation_method:
+                    description:
+                        - public ip allocation method.
+                load_balancer_backend_address_pools:
+                    description:
+                        - List of an existing load-balancer backend address pool id to associate with the network interface.
+                primary:
+                    description:
+                        - Whether the ip configuration is the primary one in the list.
+                application_security_groups:
+                    description:
+                        - List of Application security groups.
+                    sample: /subscriptions/<subsid>/resourceGroups/<rg>/providers/Microsoft.Network/applicationSecurityGroups/myASG
+        enable_accelerated_networking:
+            description:
+                - Specifies whether the network interface should be created with the accelerated networking feature or not
+        create_with_security_group:
+            description:
+                - Specifies whether a default security group should be be created with the NIC. Only applies when creating a new NIC.
+            type: bool
+        security_group:
+            description:
+                - A security group resource ID with which to associate the network interface.
+        enable_ip_forwarding:
+            description:
+                - Whether to enable IP forwarding
+        dns_servers:
+            description:
+                - Which DNS servers should the NIC lookup
+                - List of IP's
+        mac_address:
+            description:
+                - The MAC address of the network interface.
+        provisioning_state:
+            description:
+                - The provisioning state of the network interface.
+        dns_settings:
+            description:
+                - The DNS settings in network interface.
+            contains:
+                dns_servers:
+                    description: List of DNS servers IP addresses.
+                applied_dns_servers:
+                    description:
+                        - If the VM that uses this NIC is part of an Availability Set, then this list will have the union of all DNS servers
+                          from all NICs that are part of the Availability Set. This property is what is configured on each of those VMs.
+                internal_dns_name_label:
+                    description: Relative DNS name for this NIC used for internal communications between VMs in the same virtual network.
+                internal_fqdn:
+                    description: Fully qualified DNS name supporting internal communications between VMs in the same virtual network.
+'''  # NOQA
+try:
+    from msrestazure.azure_exceptions import CloudError
+    from azure.common import AzureMissingResourceHttpError, AzureHttpError
+except Exception:
+    # This is handled in azure_rm_common
+    pass
+
+from ansible.module_utils.azure_rm_common import AzureRMModuleBase, azure_id_to_dict
+
+
+AZURE_OBJECT_CLASS = 'NetworkInterface'
+
+
+class AzureRMVirtualMachineScaleSetNetworkInterfaceFacts(AzureRMModuleBase):
+
+    def __init__(self):
+
+        self.module_arg_spec = dict(
+            vmss_name=dict(type='str', required=True),
+            resource_group=dict(type='str', required=True),
+            tags=dict(type='list')
+        )
+
+        self.results = dict(
+            changed=False,
+            networkinterfaces=[]
+        )
+
+        self.vmss_name = None
+        self.resource_group = None
+        self.tags = None
+
+        super(AzureRMVirtualMachineScaleSetNetworkInterfaceFacts,
+              self).__init__(self.module_arg_spec,
+                             supports_tags=False,
+                             facts_module=True
+                             )
+
+    def exec_module(self, **kwargs):
+
+        for key in self.module_arg_spec:
+            setattr(self, key, kwargs[key])
+
+        if not (self.vmss_name and self.resource_group):
+            self.fail("Parameter error: resource group and name are required parameters.")
+
+        networkinterfaces = []
+
+        networkinterfaces = self.list_nics()
+
+        self.results['networkinterfaces'] = self._promote_properties(self.serialize_nics(networkinterfaces))
+        return self.results
+
+    def list_nics(self):
+        self.log('Get properties for {0}'.format(self.vmss_name))
+        item = None
+        try:
+            return self.network_client.network_interfaces.list_virtual_machine_scale_set_network_interfaces(self.resource_group, self.vmss_name)
+        except Exception as exc:
+            self.fail("Error fetching nic list - {0}".format(str(exc)))
+
+
+    def serialize_nics(self, raws):
+        return [self.serialize_obj(item, AZURE_OBJECT_CLASS) for item in raws] if raws else []
+
+
+    def _promote_properties(self, obj):
+        if isinstance(obj, list):
+            return [self._promote_properties(i) for i in obj]
+
+        if isinstance(obj, dict):
+            for k,v in six.iteritems(obj):
+                if isinstance(v, (list, dict)):
+                    obj.update({k: self._promote_properties(v)})
+
+            if isinstance(obj.get("properties"), dict):
+                properties = obj.pop("properties")
+                obj.update(properties)
+
+        return obj
+
+
+def main():
+    AzureRMVirtualMachineScaleSetNetworkInterfaceFacts()
+
+
+if __name__ == '__main__':
+    main()
diff --git a/ansible/roles/azure/tasks/assign_msi_single_vmss.yml b/ansible/roles/azure/tasks/assign_msi_single_vmss.yml
new file mode 100644
index 0000000..e9530f4
--- /dev/null
+++ b/ansible/roles/azure/tasks/assign_msi_single_vmss.yml
@@ -0,0 +1,30 @@
+---
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+- name: Assign User assigned Identity to Single VMSS
+  azure_rm_resource:
+    resource_group: "{{ resource_group }}"
+    provider: Compute
+    resource_type: virtualMachineScaleSets
+    resource_name: "{{ vmss_name }}"
+    api_version: '2019-03-01'
+    body:
+     location: "{{ location }}"
+     identity:
+       type: UserAssigned
+       userAssignedIdentities: "{{ UserAssignedIdentityArr|join('') }}"
diff --git a/ansible/roles/azure/tasks/create_adlsgen2.yml b/ansible/roles/azure/tasks/create_adlsgen2.yml
index 88f70be..5e7138d 100644
--- a/ansible/roles/azure/tasks/create_adlsgen2.yml
+++ b/ansible/roles/azure/tasks/create_adlsgen2.yml
@@ -223,15 +223,3 @@
     regexp: '^azure_client_id\s*=\s*|^[#]azure_client_id\s*=\s*'
     line: "azure_client_id = {{ UserAssignedIdentityInfo.response|map(attribute='properties')|map(attribute='clientId')|list|join('') }}"
 
-- name: Assign User Assigned Identity to VMSS
-  azure_rm_resource:
-    resource_group: "{{ resource_group }}"
-    provider: Compute
-    resource_type: virtualMachineScaleSets
-    resource_name: "{{ vmss_name }}"
-    api_version: '2019-03-01'
-    body:
-     location: "{{ location }}"
-     identity:
-       type: UserAssigned
-       userAssignedIdentities: "{{ UserAssignedIdentityArr|join('') }}"
diff --git a/ansible/roles/azure/tasks/create_common_resources.yml b/ansible/roles/azure/tasks/create_common_resources.yml
new file mode 100644
index 0000000..832b7a9
--- /dev/null
+++ b/ansible/roles/azure/tasks/create_common_resources.yml
@@ -0,0 +1,65 @@
+---
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#
+# These Ansible tasks only run on the client machine where Muchos runs
+# At a high level, the various sections in this file do the following:
+# 1. Create (if not already existing): an Azure resource group, virtual network / subnet
+# 2. Optionally (if the user specified) create a VM and related resources to use as a proxy host
+# 3. Create the Azure VMSS to support the nodes for use with Muchos
+# 4. Automatically populate the hosts file and associated [nodes] section in muchos.props
+#
+# SECTION 1: Create Azure RG, network and subnet
+#
+- name: Create a resource group
+  azure_rm_resourcegroup:
+    name: "{{ resource_group }}"
+    location: "{{ location }}"
+    tags:
+        deployment_type: muchos
+        application: accumulo
+
+- name: Ensure that the resource group is available
+  azure_rm_resourcegroup_info:
+    name: "{{ resource_group }}"
+  retries: 5
+  delay: 3
+  register: resource_group_facts
+  until: resource_group_facts.resourcegroups is defined and resource_group_facts.resourcegroups|map(attribute='name')|join('') is defined and resource_group_facts.resourcegroups|map(attribute='properties')|map(attribute='provisioningState')|join('') == "Succeeded"
+
+- debug:
+   msg: "Name of resource group created: {{ resource_group_facts.resourcegroups|map(attribute='name')|join('') }}"
+
+- debug:
+   msg: "Provisioning state of {{ resource_group }} is {{ resource_group_facts.resourcegroups|map(attribute='properties')|map(attribute='provisioningState')|join('') }}"
+
+- name: Create a virtual network
+  azure_rm_virtualnetwork:
+    resource_group: "{{ resource_group }}"
+    name: "{{ vnet }}"
+    address_prefixes_cidr:
+        - "{{ vnet_cidr }}"
+    tags:
+        deployment_type: muchos
+        application: accumulo
+
+- name: Create a subnet
+  azure_rm_subnet:
+    resource_group: "{{ resource_group }}"
+    virtual_network_name: "{{ vnet }}"
+    name: "{{ subnet }}"
+    address_prefix_cidr: "{{ subnet_cidr }}"
diff --git a/ansible/roles/azure/tasks/create_optional_proxy.yml b/ansible/roles/azure/tasks/create_optional_proxy.yml
new file mode 100644
index 0000000..41e484b
--- /dev/null
+++ b/ansible/roles/azure/tasks/create_optional_proxy.yml
@@ -0,0 +1,86 @@
+---
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#
+# These Ansible tasks only run on the client machine where Muchos runs
+# At a high level, the various sections in this file do the following:
+# 1. Create (if not already existing): an Azure resource group, virtual network / subnet
+# 2. Optionally (if the user specified) create a VM and related resources to use as a proxy host
+# 3. Create the Azure VMSS to support the nodes for use with Muchos
+# 4. Automatically populate the hosts file and associated [nodes] section in muchos.props
+#
+# SECTION 2: Optionally create a VM with a public IP which can act as a proxy host
+#
+- name: Create public IP address
+  azure_rm_publicipaddress:
+    resource_group: "{{ resource_group }}"
+    allocation_method: Static
+    name: "{{ azure_proxy_host }}-ip"
+  register: azure_proxy_public_ip
+  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
+
+- name: Create Network Security Group that allows SSH
+  azure_rm_securitygroup:
+    resource_group: "{{ resource_group }}"
+    name: "{{ azure_proxy_host }}-nsg"
+    rules:
+      - name: SSH
+        protocol: Tcp
+        destination_port_range: 22
+        access: Allow
+        priority: 1001
+        direction: Inbound
+  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
+
+- name: Create NIC
+  azure_rm_networkinterface:
+    resource_group: "{{ resource_group }}"
+    name: "{{ azure_proxy_host }}-nic"
+    virtual_network: "{{ vnet }}"
+    subnet_name: "{{ subnet }}"
+    ip_configurations:
+      - name: default
+        public_ip_address_name: "{{ azure_proxy_host }}-ip"
+        primary: True
+    security_group: "{{ azure_proxy_host }}-nsg"
+    enable_accelerated_networking: "{{ True if azure_proxy_host_vm_sku in accnet_capable_skus else False }}"
+  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
+
+- name: Create azure proxy virtual machine
+  azure_rm_virtualmachine:
+    resource_group: "{{ resource_group }}"
+    name: "{{ azure_proxy_host }}"
+    network_interface_names:
+      - "{{ azure_proxy_host }}-nic"
+    vm_size: "{{ azure_proxy_host_vm_sku }}"
+    admin_username: "{{ cluster_user }}"
+    ssh_password_enabled: false
+    ssh_public_keys:
+      - path: /home/{{ cluster_user }}/.ssh/authorized_keys
+        key_data: "{{ lookup('file', '~/.ssh/id_rsa.pub') }}"
+    os_disk_caching: ReadWrite
+    image:
+      offer: CentOS
+      publisher: OpenLogic
+      sku: 7.5
+      version: latest
+    managed_disk_type: "{{ managed_disk_type }}"
+    data_disks:
+     - lun: 0
+       disk_size_gb: 64
+       managed_disk_type: "{{ managed_disk_type }}"
+  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
diff --git a/ansible/roles/azure/tasks/create_vmss.yml b/ansible/roles/azure/tasks/create_vmss.yml
index 9b1438b..a2a2be8 100644
--- a/ansible/roles/azure/tasks/create_vmss.yml
+++ b/ansible/roles/azure/tasks/create_vmss.yml
@@ -24,93 +24,11 @@
 # 3. Create the Azure VMSS to support the nodes for use with Muchos
 # 4. Automatically populate the hosts file and associated [nodes] section in muchos.props
 
-# SECTION 1: Create Azure RG, network and subnet
-- name: Create a resource group
-  azure_rm_resourcegroup:
-    name: "{{ resource_group }}"
-    location: "{{ location }}"
-    tags:
-        deployment_type: muchos
-        application: accumulo
-
-- name: Create a virtual network
-  azure_rm_virtualnetwork:
-    resource_group: "{{ resource_group }}"
-    name: "{{ vnet }}"
-    address_prefixes_cidr:
-        - "{{ vnet_cidr }}"
-    tags:
-        deployment_type: muchos
-        application: accumulo
-
-- name: Create a subnet
-  azure_rm_subnet:
-    resource_group: "{{ resource_group }}"
-    virtual_network_name: "{{ vnet }}"
-    name: "{{ subnet }}"
-    address_prefix_cidr: "{{ subnet_cidr }}"
-
-# SECTION 2: Optionally create a VM with a public IP which can act as a proxy host
-- name: Create public IP address
-  azure_rm_publicipaddress:
-    resource_group: "{{ resource_group }}"
-    allocation_method: Static
-    name: "{{ azure_proxy_host }}-ip"
-  register: azure_proxy_public_ip
-  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
-
-- name: Create Network Security Group that allows SSH
-  azure_rm_securitygroup:
-    resource_group: "{{ resource_group }}"
-    name: "{{ azure_proxy_host }}-nsg"
-    rules:
-      - name: SSH
-        protocol: Tcp
-        destination_port_range: 22
-        access: Allow
-        priority: 1001
-        direction: Inbound
-  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
-
-- name: Create NIC
-  azure_rm_networkinterface:
-    resource_group: "{{ resource_group }}"
-    name: "{{ azure_proxy_host }}-nic"
-    virtual_network: "{{ vnet }}"
-    subnet: "{{ subnet }}"
-    ip_configurations:
-      - name: default
-        public_ip_address_name: "{{ azure_proxy_host }}-ip"
-        primary: True
-    security_group: "{{ azure_proxy_host }}-nsg"
-  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
-
-- name: Create azure proxy virtual machine
-  azure_rm_virtualmachine:
-    resource_group: "{{ resource_group }}"
-    name: "{{ azure_proxy_host }}"
-    network_interface_names:
-      - "{{ azure_proxy_host }}-nic"
-    vm_size: Standard_D8s_v3
-    admin_username: "{{ admin_username }}"
-    ssh_password_enabled: false
-    ssh_public_keys:
-      - path: /home/{{ admin_username }}/.ssh/authorized_keys
-        key_data: "{{ lookup('file', '~/.ssh/id_rsa.pub') }}"
-    os_disk_caching: ReadWrite
-    managed_disk_type: Standard_LRS
-    image:
-      offer: CentOS
-      publisher: OpenLogic
-      sku: 7.5
-      version: latest
-    data_disks:
-     - lun: 0
-       disk_size_gb: 64
-       managed_disk_type: Standard_LRS
-  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
+# Moved SECTION1 to create_common_resources.yml
+# Moved SECTION2 to create_optional_proxy.yml
 
 # SECTION 3: Create the Azure VMSS for the nodes used by Muchos
+
 - name: Create luns dictionary
   set_fact:
     luns_dict: "{{ luns_dict | default ([]) + [{ 'lun': item, 'disk_size_gb': disk_size_gb , 'caching': None } ] }}"
@@ -139,7 +57,7 @@
     tier: Standard
     managed_disk_type: "{{ managed_disk_type }}"
     os_disk_caching: ReadWrite
-    enable_accelerated_networking: yes
+    enable_accelerated_networking: "{{ accnet_capable }}"
     single_placement_group: "{{ single_placement_group | default(omit) }}"
     image:
       offer: "{{ image_offer if image_offer else omit }}"
@@ -152,28 +70,54 @@
   - image_publisher: "{{ azure_image_reference.split('|')[1] }}"
   - image_sku: "{{ azure_image_reference.split('|')[2] }}"
   - image_version: "{{ azure_image_reference.split('|')[3] }}"
+  - accnet_capable: "{{ True if vm_sku in accnet_capable_skus else False }}"
   tags: create_vmss
 
 # SECTION 4: Automatically populate entries in the hosts file and in the muchos.props file, based on the VMSS node details
+# Refactored below tasks to use Azure modules
 - name: get instance ids
-  shell: "set -o pipefail && az vmss list-instances --resource-group {{ resource_group }} --name {{ vmss_name }} --query '[].[{name: name,vmresourceid: id}]' -o json | sed -e 's/{{ vmss_name }}_/{{ vmss_name }}-/g'"
-  args:
-    executable: bash
-  register: instance_list
-- name: set instance_list
-  set_fact:
-    instances: "{{ instance_list.stdout|from_json }}"
+  azure_rm_virtualmachinescalesetinstance_info:
+    resource_group: "{{ resource_group }}"
+    vmss_name: "{{ vmss_name }}"
+  register: vmss_instances
 
-- name: get private ip addresses
-  command: "az vmss nic list --resource-group {{ resource_group }} --vmss-name {{ vmss_name }} --query '[].[{ip: ipConfigurations[0].privateIpAddress, vmresourceid: virtualMachine.id}]' -o json"
-  register: ip_list
-- name: obtain ip address list
+- name: replace underscores in instance names with dashes
   set_fact:
-    addresses: "{{ ip_list.stdout|from_json }}"
-- name: create instance name and ip address dictionaries
+    hostnames: "{{ vmss_instances.instances | map(attribute='name') | map('replace', '_', '-') | list }}"
+
+- name: Ensures host_vars sub-dir exists
+  file:
+     path: "{{ deploy_path }}/ansible/host_vars/"
+     state: directory
+     recurse: yes
+
+- name: Get VMSS nic list
+  azure_rm_virtualmachinescaleset_nic_list_facts:
+    resource_group: "{{ resource_group }}"
+    vmss_name: "{{ vmss_name }}"
+  register: vmss_nic_list
+
+- name: Get VM hostname to IP mapping
   set_fact:
-    instances_dict: "{{ instances | json_query('[].{key: vmresourceid, value: name}') }}"
-    addresses_dict: "{{ addresses | json_query('[].{key: vmresourceid, value: ip}') }}"
+    hostname_ip_pairs: |
+      {%- set vmname_ips = [] -%}
+      {%- if azure_proxy_host is defined and azure_proxy_host -%}
+      {%- set _ = vmname_ips.append({'name': azure_proxy_host, 'ip': azure_proxy_public_ip.state.ip_address }) -%}
+      {%- endif -%}
+      {%- set vmid_names = {} -%}
+      {%- for instance in vmss_instances.instances -%}
+      {%- set _ = vmid_names.__setitem__(instance.id, instance.name.replace('_','-')) -%}
+      {%- endfor -%}
+      {%- set vmid_ips = {} -%}
+      {%- for interface in vmss_nic_list.networkinterfaces -%}
+      {%- if interface.virtualMachine is defined -%}
+      {%- set _ = vmid_ips.__setitem__(interface.virtualMachine.id, interface.ipConfigurations[0].privateIPAddress) -%}
+      {%- endif -%}
+      {%- endfor -%}
+      {%- for vmid in vmid_names -%}
+      {%- set _ = vmname_ips.append({'name': vmid_names[vmid], 'ip': vmid_ips[vmid]}) -%}
+      {%- endfor -%}
+      {{ vmname_ips }}
 
 - name: Ensures hosts sub-dir exists
   file:
@@ -181,29 +125,12 @@
      state: directory
      recurse: yes
 
-- name: Ensure empty hosts file is created
-  copy:
-    content: ""
+- name: Write hosts file
+  template:
+    src: hostname_ip_mappings.j2
     dest: "{{ deploy_path }}/conf/hosts/{{ vmss_name }}"
-    force: yes
     mode: 0644
 
-- name: Write azure proxy to hosts file
-  lineinfile:
-    path: "{{ deploy_path }}/conf/hosts/{{ vmss_name }}"
-    line: "{{ azure_proxy_host }} {{ azure_proxy_public_ip.state.ip_address }}"
-  when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
-
-- name: Write VMSS vms to hosts file
-  lineinfile:
-    path: "{{ deploy_path }}/conf/hosts/{{ vmss_name }}"
-    line: "{{ item1[0] }} {{ item2[0] }}"
-  with_items: "{{ instances_dict | json_query('[].key') }}"
-  vars:
-    myquery: "[?key=='{{ item }}'].value"
-    item1: "{{ instances_dict | json_query(myquery) }}"
-    item2: "{{ addresses_dict | json_query(myquery)  }}"
-
 - name: Clear section
   ini_file:
     path: "{{ deploy_path }}/conf/muchos.props"
@@ -224,53 +151,40 @@
     line: "{{ azure_proxy_host }} = client"
   when: azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None
 
-- name: Assign Accumulo master, HDFS components to the first node of the cluster
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = namenode,resourcemanager,accumulomaster,zookeeper"
-  with_items: "{{ instances_dict | json_query('[0].value') }}"
+- block:
+  - name: Assign Accumulo master, HDFS components to the first node of the cluster
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ hostnames[0] }} = namenode,resourcemanager,accumulomaster,zookeeper"
+  - name: Assign metrics to the second node of the cluster
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ hostnames[1] }} = metrics"
+  - name: Add worker nodes to muchos.props
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ item }} = worker"
+    with_items: "{{ hostnames | json_query('[2:]') }}"
   when: not hdfs_ha
 
-- name: Assign metrics to the second node of the cluster
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = metrics"
-  with_items: "{{ instances_dict | json_query('[1].value') }}"
-  when: not hdfs_ha
-
-- name: Add worker nodes to muchos.props
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = worker"
-  with_items: "{{ instances_dict | json_query('[2:].value') }}"
-  when: not hdfs_ha
-
-- name: Assign Accumulo master, HDFS HA components cluster roles to the first node of the cluster
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = namenode,resourcemanager,accumulomaster,zookeeper,journalnode,zkfc"
-  with_items: "{{ instances_dict | json_query('[0].value') }}"
-  when: hdfs_ha
-
-- name: Assign Accumulo master, HDFS HA components cluster roles to the second node of the cluster
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = zookeeper,metrics,journalnode,namenode,zkfc,accumulomaster,resourcemanager"
-  with_items: "{{ instances_dict | json_query('[1].value') }}"
-  when: hdfs_ha
-
-- name: Assign HDFS HA components cluster roles to the third node of the cluster
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = journalnode,zookeeper"
-  with_items: "{{ instances_dict | json_query('[2].value') }}"
-  when: hdfs_ha
-
-- name: Add worker nodes to muchos.props
-  lineinfile:
-    path: "{{ deploy_path }}/conf/muchos.props"
-    line: "{{ item }} = worker"
-  with_items: "{{ instances_dict | json_query('[3:].value') }}"
+- block:
+  - name: Assign Accumulo master, HDFS HA components cluster roles to the first node of the cluster
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ hostnames[0] }} = namenode,resourcemanager,accumulomaster,zookeeper,journalnode,zkfc"
+  - name: Assign Accumulo master, HDFS HA components cluster roles to the second node of the cluster
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ hostnames[1] }} = zookeeper,metrics,journalnode,namenode,zkfc,accumulomaster,resourcemanager"
+  - name: Assign HDFS HA components cluster roles to the third node of the cluster
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ hostnames[2] }} = journalnode,zookeeper"
+  - name: Add worker nodes to muchos.props
+    lineinfile:
+      path: "{{ deploy_path }}/conf/muchos.props"
+      line: "{{ item }} = worker"
+    with_items: "{{ hostnames | json_query('[3:]') }}"
   when: hdfs_ha
 
 - name: Change proxy hostname to azure proxy host in muchos.props
@@ -284,6 +198,5 @@
   lineinfile:
     path: "{{ deploy_path }}/conf/muchos.props"
     regexp: '^proxy_hostname\s*=\s*'
-    line: "proxy_hostname = {{ item }}"
-  with_items: "{{ instances_dict | json_query('[0].value') }}"
+    line: "proxy_hostname = {{ hostnames[0] }}"
   when: not (azure_proxy_host is defined and azure_proxy_host and azure_proxy_host != None)
diff --git a/ansible/roles/azure/tasks/main.yml b/ansible/roles/azure/tasks/main.yml
index f8465a8..d252833 100644
--- a/ansible/roles/azure/tasks/main.yml
+++ b/ansible/roles/azure/tasks/main.yml
@@ -25,8 +25,12 @@
     Update it using pip install --upgrade ansible[azure]==2.9.13"
   when: use_adlsg2
 
-- import_tasks: create_vmss.yml
+- import_tasks: create_common_resources.yml
 - import_tasks: create_adlsgen2.yml
   when: use_adlsg2
+- import_tasks: create_optional_proxy.yml
+- import_tasks: create_vmss.yml
+- import_tasks: assign_msi_single_vmss.yml
+  when: use_adlsg2
 - import_tasks: create_log_analytics_ws.yml
   when: az_oms_integration_needed and (az_logs_id is not defined or (not az_logs_id) or az_logs_id == None)
diff --git a/ansible/roles/azure/templates/hostname_ip_mappings.j2 b/ansible/roles/azure/templates/hostname_ip_mappings.j2
new file mode 100644
index 0000000..e17ba5a
--- /dev/null
+++ b/ansible/roles/azure/templates/hostname_ip_mappings.j2
@@ -0,0 +1,3 @@
+{%- for name_ip_pair in hostname_ip_pairs %}
+{{ name_ip_pair['name']}} {{ name_ip_pair['ip'] }}
+{% endfor -%}
diff --git a/conf/checksums b/conf/checksums
index 0a8b8e0..9dd58e6 100644
--- a/conf/checksums
+++ b/conf/checksums
@@ -39,6 +39,7 @@
 spark:2.3.3:a5c4a54db5163df25110534f36e13d6ae65758d45d1bc8bfcd7e27304e716a23
 spark:2.3.2:3387107155d62f04ccf6bcaf2e00a69a0de5ae5df875348d93147743c206f0a8
 spark:2.2.2:023b2fea378b3dd0fee2d5d1de6bfaf2d8349aefe7be97a9cbcf03bbacc428d7
+zookeeper:3.6.3:3f7b1b7d9cf5647d52ad0076c922e108fa956e986b5624667c493cf6d8ff09d3ca88f623c79a799fe49c72e868cb3c9d0f77cb69608de74a183b2cbad10bc827
 zookeeper:3.6.2:caff5111bb6876b7124760bc006e6fa2523efa54b99321a3c9cd8192ea0d5596abc7d70a054b1aac9b20a411407dae7611c7aba870c23bff28eb1643ba499199
 zookeeper:3.6.1:1c5cb4d9886fae41bf244a446dd874b73c0fff7a5fc2dda4305041964420cde21e59b388dfd2551464a46bb6918d9d3c3c01ebe68fdbe782074ee360aa830c7d
 zookeeper:3.5.9:0e5a64713abc6f36d961dd61a06f681868171a9d9228366e512a01324806d263e05508029c94d8e18307811867cdc39d848e736c252bf56c461273ef74c66a45
diff --git a/conf/muchos.props.example b/conf/muchos.props.example
index c6f0056..54fcefb 100644
--- a/conf/muchos.props.example
+++ b/conf/muchos.props.example
@@ -37,7 +37,7 @@
 accumulo_password = secret
 # Software versions (make sure you have a corresponding entry for the checksum in conf/checksums)
 hadoop_version = 3.2.1
-zookeeper_version = 3.6.2
+zookeeper_version = 3.6.3
 spark_version = 2.4.5
 fluo_version = 1.2.0
 fluo_yarn_version = 1.0.0