From 31f793e3640c4e10abc312fc87ed798dcb1056a2 Mon Sep 17 00:00:00 2001
From: Travis Prescott <trpresco@microsoft.com>
Date: Wed, 15 Mar 2017 10:08:25 -0700
Subject: [PATCH 1/2] Enable separate os/disk caching mechanisms.

---
 .../azure/cli/command_modules/vm/_params.py   | 13 +++---
 .../command_modules/vm/_template_builder.py   | 25 ++++++------
 .../cli/command_modules/vm/_validators.py     |  6 +--
 .../azure/cli/command_modules/vm/custom.py    | 40 +++++++++----------
 .../azure-cli-vm/tests/test_vm_commands.py    |  2 +-
 5 files changed, 44 insertions(+), 42 deletions(-)

diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
index 58f5df48ce7..52ddd8eda06 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
@@ -68,9 +68,9 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
 register_cli_argument('vm unmanaged-disk detach', 'disk_name', options_list=('--name', '-n'), help='The data disk name.')
 register_cli_argument('vm unmanaged-disk', 'disk_size', help='Size of disk (GiB)', default=1023, type=int)
 register_cli_argument('vm unmanaged-disk', 'new', action="store_true", help='create a new disk')
-register_cli_argument('vm unmanaged-disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virutal Machine size.')
+register_cli_argument('vm unmanaged-disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virtual Machine size.')
 register_cli_argument('vm unmanaged-disk', 'vhd_uri', help="virtual hard disk's uri. For example:https://mystorage.blob.core.windows.net/vhds/d1.vhd")
-register_cli_argument('vm unmanaged-disk', 'caching', help='Host caching policy', default=CachingTypes.none.value, **enum_choice_list(CachingTypes))
+register_cli_argument('vm', 'caching', help='Disk caching policy', **enum_choice_list(CachingTypes))
 
 for item in ['attach', 'detach']:
     register_cli_argument('vm unmanaged-disk {}'.format(item), 'vm_name', arg_type=existing_vm_name, options_list=('--vm-name',), id_part=None)
@@ -83,7 +83,7 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
 register_cli_argument('vm disk', 'new', action="store_true", help='create a new disk')
 register_cli_argument('vm disk', 'sku', arg_type=disk_sku)
 register_cli_argument('vm disk', 'size_gb', options_list=('--size-gb', '-z'), help='size in GB.')
-register_cli_argument('vm disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virutal Machine size.')
+register_cli_argument('vm disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virtual Machine size.')
 
 register_cli_argument('vm availability-set', 'availability_set_name', name_arg_type, id_part='name',
                       completer=get_resource_name_completion_list('Microsoft.Compute/availabilitySets'), help='Name of the availability set')
@@ -122,8 +122,9 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
 register_cli_argument('vmss', 'instance_id', id_part='child_name')
 register_cli_argument('vmss', 'instance_ids', multi_ids_type, help='Space separated list of IDs (ex: 1 2 3 ...) or * for all instances. If not provided, the action will be applied on the scaleset itself')
 register_cli_argument('vmss', 'tags', tags_type)
+register_cli_argument('vmss', 'caching', help='Disk caching policy', **enum_choice_list(CachingTypes))
 
-register_cli_argument('vmss disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virutal Machine instance size.')
+register_cli_argument('vmss disk', 'lun', type=int, help='0-based logical unit number (LUN). Max value depends on the Virtual Machine instance size.')
 register_cli_argument('vmss disk', 'size_gb', options_list=('--size-gb', '-z'), help='size in GB.')
 register_cli_argument('vmss disk', 'vmss_name', vmss_name_type, completer=get_resource_name_completion_list('Microsoft.Compute/virtualMachineScaleSets'))
 
@@ -222,11 +223,12 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
     register_cli_argument(scope, 'public_ip_address_allocation', help=None, arg_group='Network', **enum_choice_list(['dynamic', 'static']))
     register_cli_argument(scope, 'public_ip_address_dns_name', help='Globally unique DNS name for a newly created Public IP.', arg_group='Network')
     register_cli_argument(scope, 'secrets', multi_ids_type, help='One or many Key Vault secrets as JSON strings or files via \'@<file path>\' containing \'[{ "sourceVault": { "id": "value" }, "vaultCertificates": [{ "certificateUrl": "value", "certificateStore": "cert store name (only on windows)"}] }]\'', type=file_type, completer=FilesCompleter())
+    register_cli_argument(scope, 'os_caching', options_list=['--storage-caching', '--os-caching'], arg_group='Storage', help='Storage caching type for the VM OS disk.', **enum_choice_list(CachingTypes))
+    register_cli_argument(scope, 'data_caching', options_list=['--data-caching'], arg_group='Storage', help='Storage caching type for the VM data disk(s).', **enum_choice_list(CachingTypes))
 
 register_cli_argument('vm create', 'vm_name', name_arg_type, id_part=None, help='Name of the virtual machine.', validator=process_vm_create_namespace, completer=None)
 register_cli_argument('vm create', 'attach_os_disk', help='Attach an existing OS disk to the VM. Can use the name or ID of a managed disk or the URI to an unmanaged disk VHD.')
 register_cli_argument('vm create', 'availability_set', help='Name or ID of an existing availability set to add the VM to. None by default.')
-register_cli_argument('vm create', 'storage_caching', help='Storage caching type for the VM OS disk. Default: ReadWrite.', arg_group='Storage', **enum_choice_list(['ReadWrite', 'ReadOnly']))
 
 register_cli_argument('vmss create', 'vmss_name', name_arg_type, id_part=None, help='Name of the virtual machine scale set.', validator=process_vmss_create_namespace)
 register_cli_argument('vmss create', 'load_balancer', help='Name to use when creating a new load balancer (default) or referencing an existing one. Can also reference an existing load balancer by ID or specify "" for none.', arg_group='Load Balancer')
@@ -237,7 +239,6 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
 register_cli_argument('vmss create', 'disable_overprovision', help='Overprovision option (see https://azure.microsoft.com/en-us/documentation/articles/virtual-machine-scale-sets-overview/ for details).', action='store_true')
 register_cli_argument('vmss create', 'upgrade_policy_mode', help=None, **enum_choice_list(UpgradeMode))
 register_cli_argument('vmss create', 'vm_sku', help='Size of VMs in the scale set.  See https://azure.microsoft.com/en-us/pricing/details/virtual-machines/ for size info.')
-register_cli_argument('vmss create', 'storage_caching', help='Storage caching type for the VM OS disk. Default: ReadOnly.', arg_group='Storage', **enum_choice_list(['ReadWrite', 'ReadOnly']))
 
 register_cli_argument('vm encryption', 'volume_type', help='Type of volume that the encryption operation is performed on', **enum_choice_list(['DATA', 'OS', 'ALL']))
 register_cli_argument('vm encryption', 'force', action='store_true', help='continue with encryption operations regardless of the warnings')
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
index 821146e3440..68610f4df87 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
@@ -253,7 +253,7 @@ def build_vm_resource(  # pylint: disable=too-many-locals
         name, location, tags, size, storage_profile, nics, admin_username,
         availability_set_id=None, admin_password=None, ssh_key_value=None, ssh_key_path=None,
         image_reference=None, os_disk_name=None, custom_image_os_type=None,
-        storage_caching=None, storage_sku=None,
+        os_caching=None, data_caching=None, storage_sku=None,
         os_publisher=None, os_offer=None, os_sku=None, os_version=None, os_vhd_uri=None,
         attach_os_disk=None, data_disk_sizes_gb=None, image_data_disks=None,
         custom_data=None, secrets=None):
@@ -296,7 +296,7 @@ def _build_storage_profile():
                 'osDisk': {
                     'createOption': 'fromImage',
                     'name': os_disk_name,
-                    'caching': storage_caching,
+                    'caching': os_caching,
                     'osType': custom_image_os_type,
                     'image': {'uri': image_reference},
                     'vhd': {'uri': os_vhd_uri}
@@ -306,7 +306,7 @@ def _build_storage_profile():
                 'osDisk': {
                     'createOption': 'fromImage',
                     'name': os_disk_name,
-                    'caching': storage_caching,
+                    'caching': os_caching,
                     'vhd': {'uri': os_vhd_uri}
                 },
                 'imageReference': {
@@ -328,7 +328,7 @@ def _build_storage_profile():
                 'osDisk': {
                     'createOption': 'fromImage',
                     'name': os_disk_name,
-                    'caching': storage_caching,
+                    'caching': os_caching,
                     'managedDisk': {'storageAccountType': storage_sku}
                 },
                 'imageReference': {
@@ -342,7 +342,7 @@ def _build_storage_profile():
                 'osDisk': {
                     'createOption': 'fromImage',
                     'name': os_disk_name,
-                    'caching': storage_caching,
+                    'caching': os_caching,
                     'managedDisk': {'storageAccountType': storage_sku}
                 },
                 "imageReference": {
@@ -361,7 +361,7 @@ def _build_storage_profile():
         }
         profile = storage_profiles[storage_profile.name]
         return _build_data_disks(profile, data_disk_sizes_gb, image_data_disks,
-                                 storage_caching, storage_sku)
+                                 data_caching, storage_sku)
 
     vm_properties = {
         'hardwareProfile': {'vmSize': size},
@@ -389,7 +389,7 @@ def _build_storage_profile():
 
 
 def _build_data_disks(profile, data_disk_sizes_gb, image_data_disks,
-                      storage_caching, storage_sku):
+                      data_caching, storage_sku):
     if data_disk_sizes_gb is not None:
         profile['dataDisks'] = []
         for image_data_disk in image_data_disks or []:
@@ -403,7 +403,7 @@ def _build_data_disks(profile, data_disk_sizes_gb, image_data_disks,
                 'lun': lun,
                 'createOption': "empty",
                 'diskSizeGB': int(size),
-                'caching': storage_caching,
+                'caching': data_caching,
                 'managedDisk': {'storageAccountType': storage_sku}
             })
             lun = lun + 1
@@ -498,7 +498,8 @@ def build_vmss_resource(name, naming_prefix, location, tags, overprovision, upgr
                         vm_sku, instance_count, ip_config_name, nic_name, subnet_id,
                         admin_username, authentication_type,
                         storage_profile, os_disk_name,
-                        storage_caching, storage_sku, data_disk_sizes_gb, image_data_disks, os_type,
+                        os_caching, data_caching, storage_sku, data_disk_sizes_gb,
+                        image_data_disks, os_type,
                         image=None, admin_password=None, ssh_key_value=None, ssh_key_path=None,
                         os_publisher=None, os_offer=None, os_sku=None, os_version=None,
                         backend_address_pool_id=None, inbound_nat_pool_id=None,
@@ -526,7 +527,7 @@ def build_vmss_resource(name, naming_prefix, location, tags, overprovision, upgr
     if storage_profile in [StorageProfile.SACustomImage, StorageProfile.SAPirImage]:
         storage_properties['osDisk'] = {
             'name': os_disk_name,
-            'caching': storage_caching,
+            'caching': os_caching,
             'createOption': 'FromImage',
         }
 
@@ -542,7 +543,7 @@ def build_vmss_resource(name, naming_prefix, location, tags, overprovision, upgr
     elif storage_profile in [StorageProfile.ManagedPirImage, StorageProfile.ManagedPirImage]:
         storage_properties['osDisk'] = {
             'createOption': 'FromImage',
-            'caching': storage_caching,
+            'caching': os_caching,
             'managedDisk': {'storageAccountType': storage_sku}
         }
 
@@ -559,7 +560,7 @@ def build_vmss_resource(name, naming_prefix, location, tags, overprovision, upgr
         }
 
     storage_profile = _build_data_disks(storage_properties, data_disk_sizes_gb,
-                                        image_data_disks, storage_caching,
+                                        image_data_disks, data_caching,
                                         storage_sku)
 
     # Build OS Profile
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
index 75be7e57989..b39987b4902 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
@@ -188,7 +188,7 @@ def _get_storage_profile_description(profile):
 
 storage_profile_param_options = {
     'os_disk_name': '--os-disk-name',
-    'storage_caching': '--storage-caching',
+    'os_caching': '--os-caching',
     'os_type': '--os-type',
     'attach_os_disk': '--attach-os-disk',
     'image': '--image',
@@ -275,7 +275,7 @@ def _validate_vm_create_storage_profile(namespace, for_scale_set=False):
 
     elif namespace.storage_profile == StorageProfile.ManagedSpecializedOSDisk:
         required = ['os_type', 'attach_os_disk']
-        forbidden = ['os_disk_name', 'storage_caching', 'storage_account',
+        forbidden = ['os_disk_name', 'os_caching', 'storage_account',
                      'storage_container_name', 'use_unmanaged_disk', 'storage_sku']
         _validate_managed_disk_sku(namespace.storage_sku)
 
@@ -289,7 +289,7 @@ def _validate_vm_create_storage_profile(namespace, for_scale_set=False):
 
     elif namespace.storage_profile == StorageProfile.SASpecializedOSDisk:
         required = ['os_type', 'attach_os_disk', 'use_unmanaged_disk']
-        forbidden = ['os_disk_name', 'storage_caching', 'image', 'storage_account',
+        forbidden = ['os_disk_name', 'os_caching', 'image', 'storage_account',
                      'storage_container_name', 'data_disk_sizes_gb', 'storage_sku']
 
     else:
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
index 50bf12ff24d..49a7e42d8fb 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
@@ -325,7 +325,7 @@ def update_managed_disk(instance, size_gb=None, sku=None):
 
 
 def attach_managed_data_disk(resource_group_name, vm_name, disk,
-                             new=False, sku=None, size_gb=None, lun=None):
+                             new=False, sku=None, size_gb=None, lun=None, caching=None):
     '''attach a managed disk'''
     vm = get_vm(resource_group_name, vm_name)
     from azure.mgmt.compute.models import (CreationData, DiskCreateOptionTypes,
@@ -340,7 +340,7 @@ def attach_managed_data_disk(resource_group_name, vm_name, disk,
             raise CLIError('usage error: --size-gb required to create an empty disk for attach')
         data_disk = DataDisk(lun, DiskCreateOptionTypes.empty,
                              name=parse_resource_id(disk)['name'],
-                             disk_size_gb=size_gb)
+                             disk_size_gb=size_gb, caching=caching)
     else:
         params = ManagedDiskParameters(id=disk,
                                        storage_account_type=sku)
@@ -361,7 +361,8 @@ def detach_data_disk(resource_group_name, vm_name, disk_name):
     set_vm(vm)
 
 
-def attach_managed_data_disk_to_vmss(resource_group_name, vmss_name, size_gb, lun=None):
+def attach_managed_data_disk_to_vmss(resource_group_name, vmss_name, size_gb, lun=None,
+                                     caching=None):
     from azure.mgmt.compute.models import (DiskCreateOptionTypes,
                                            VirtualMachineScaleSetDataDisk)
     client = _compute_client_factory()
@@ -373,7 +374,7 @@ def attach_managed_data_disk_to_vmss(resource_group_name, vmss_name, size_gb, lu
         luns = [d.lun for d in data_disks]
         lun = max(luns) + 1 if luns else 0
     data_disk = VirtualMachineScaleSetDataDisk(lun, DiskCreateOptionTypes.empty,
-                                               disk_size_gb=size_gb)
+                                               disk_size_gb=size_gb, caching=caching)
     data_disks.append(data_disk)
     vmss.virtual_machine_profile.storage_profile.data_disks = data_disks
     return client.virtual_machine_scale_sets.create_or_update(resource_group_name, vmss_name, vmss)
@@ -1508,7 +1509,7 @@ def create_vm(vm_name, resource_group_name, image=None,
               public_ip_address=None, public_ip_address_allocation='dynamic',
               public_ip_address_dns_name=None,
               os_disk_name=None, os_type=None, storage_account=None,
-              storage_caching=None, storage_container_name=None,
+              os_caching=None, data_caching=None, storage_container_name=None,
               storage_sku=None, use_unmanaged_disk=False,
               attach_os_disk=None, data_disk_sizes_gb=None, image_data_disks=None,
               vnet_name=None, vnet_address_prefix='10.0.0.0/16',
@@ -1535,7 +1536,6 @@ def create_vm(vm_name, resource_group_name, image=None,
     tags = tags or {}
     os_disk_name = os_disk_name or 'osdisk_{}'.format(random_string(10))
     storage_container_name = storage_container_name or 'vhds'
-    storage_caching = storage_caching or 'ReadWrite'
 
     # Build up the ARM template
     master_template = ArmTemplateBuilder()
@@ -1625,7 +1625,7 @@ def create_vm(vm_name, resource_group_name, image=None,
     vm_resource = build_vm_resource(
         vm_name, location, tags, size, storage_profile, nics, admin_username, availability_set,
         admin_password, ssh_key_value, ssh_dest_key_path, image, os_disk_name,
-        os_type, storage_caching, storage_sku, os_publisher, os_offer, os_sku, os_version,
+        os_type, os_caching, data_caching, storage_sku, os_publisher, os_offer, os_sku, os_version,
         os_vhd_uri, attach_os_disk, data_disk_sizes_gb, image_data_disks, custom_data, secrets)
     vm_resource['dependsOn'] = vm_dependencies
 
@@ -1662,7 +1662,7 @@ def create_vmss(vmss_name, resource_group_name, image,
                 load_balancer=None, backend_pool_name=None, nat_pool_name=None, backend_port=None,
                 public_ip_address=None, public_ip_address_allocation='dynamic',
                 public_ip_address_dns_name=None,
-                storage_caching=None,
+                os_caching=None, data_caching=None,
                 storage_container_name=None, storage_sku=None,
                 os_type=None, os_disk_name=None,
                 use_unmanaged_disk=False, data_disk_sizes_gb=None, image_data_disks=None,
@@ -1691,7 +1691,6 @@ def create_vmss(vmss_name, resource_group_name, image,
     tags = tags or {}
     os_disk_name = os_disk_name or 'osdisk_{}'.format(random_string(10))
     storage_container_name = storage_container_name or 'vhds'
-    storage_caching = storage_caching or 'ReadOnly'
 
     # Build up the ARM template
     master_template = ArmTemplateBuilder()
@@ -1736,14 +1735,22 @@ def create_vmss(vmss_name, resource_group_name, image,
         lb_resource['dependsOn'] = lb_dependencies
         master_template.add_resource(lb_resource)
 
+    scrubbed_name = vmss_name.replace('-', '').lower()[:5]
+    naming_prefix = '{}{}'.format(scrubbed_name,
+                                  random_string(9 - len(scrubbed_name), force_lower=True))
+
     if storage_profile in [StorageProfile.SACustomImage, StorageProfile.SAPirImage]:
         master_template.add_resource(build_vmss_storage_account_pool_resource(
             'storageLoop', location, tags, storage_sku))
+        master_template.add_variable('storageAccountNames', [
+            '{}{}'.format(naming_prefix, x) for x in range(5)
+        ])
+        master_template.add_variable('vhdContainers', [
+            "[concat('https://', variables('storageAccountNames')[{}], '.blob.{}/{}')]".format(
+                x, CLOUD.suffixes.storage_endpoint, storage_container_name) for x in range(5)
+        ])
         vmss_dependencies.append('storageLoop')
 
-    scrubbed_name = vmss_name.replace('-', '').lower()[:5]
-    naming_prefix = '{}{}'.format(scrubbed_name,
-                                  random_string(9 - len(scrubbed_name), force_lower=True))
     backend_address_pool_id = None
     inbound_nat_pool_id = None
     if is_valid_resource_id(load_balancer):
@@ -1773,7 +1780,7 @@ def create_vmss(vmss_name, resource_group_name, image,
                                         vm_sku, instance_count,
                                         ip_config_name, nic_name, subnet_id, admin_username,
                                         authentication_type, storage_profile,
-                                        os_disk_name, storage_caching,
+                                        os_disk_name, os_caching, data_caching,
                                         storage_sku, data_disk_sizes_gb, image_data_disks,
                                         os_type, image, admin_password,
                                         ssh_key_value, ssh_dest_key_path,
@@ -1784,13 +1791,6 @@ def create_vmss(vmss_name, resource_group_name, image,
     vmss_resource['dependsOn'] = vmss_dependencies
 
     master_template.add_resource(vmss_resource)
-    master_template.add_variable('storageAccountNames', [
-        '{}{}'.format(naming_prefix, x) for x in range(5)
-    ])
-    master_template.add_variable('vhdContainers', [
-        "[concat('https://', variables('storageAccountNames')[{}], '.blob.{}/{}')]".format(
-            x, CLOUD.suffixes.storage_endpoint, storage_container_name) for x in range(5)
-    ])
     master_template.add_output('VMSS', vmss_name, 'Microsoft.Compute', 'virtualMachineScaleSets',
                                output_type='object')
     template = master_template.build()
diff --git a/src/command_modules/azure-cli-vm/tests/test_vm_commands.py b/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
index 237ae93652f..6388ece4213 100644
--- a/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
+++ b/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
@@ -1443,7 +1443,7 @@ def body(self):
 
         self.cmd('vmss create --image Debian --admin-password Test1234@! -l westus'
                  ' -g {} -n {} --disable-overprovision --instance-count {}'
-                 ' --storage-caching {} --upgrade-policy-mode {}'
+                 ' --os-caching {} --upgrade-policy-mode {}'
                  ' --authentication-type password --admin-username myadmin --public-ip-address {}'
                  ' --data-disk-sizes-gb 1 --vm-sku Standard_D2_v2'
                  .format(self.resource_group, vmss_name, instance_count, caching,

From 79cd649e53f984f3842eab3224c7af0040613e34 Mon Sep 17 00:00:00 2001
From: Travis Prescott <trpresco@microsoft.com>
Date: Thu, 16 Mar 2017 16:43:42 -0700
Subject: [PATCH 2/2] Code review feedback.

---
 .../azure-cli-vm/azure/cli/command_modules/vm/_params.py    | 4 ++--
 .../azure/cli/command_modules/vm/_template_builder.py       | 3 ++-
 .../azure/cli/command_modules/vm/_validators.py             | 6 +++---
 .../azure-cli-vm/azure/cli/command_modules/vm/custom.py     | 2 ++
 src/command_modules/azure-cli-vm/tests/test_vm_commands.py  | 2 +-
 5 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
index 52ddd8eda06..fb0d3ea1849 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_params.py
@@ -223,8 +223,8 @@ def get_vm_size_completion_list(prefix, action, parsed_args, **kwargs):  # pylin
     register_cli_argument(scope, 'public_ip_address_allocation', help=None, arg_group='Network', **enum_choice_list(['dynamic', 'static']))
     register_cli_argument(scope, 'public_ip_address_dns_name', help='Globally unique DNS name for a newly created Public IP.', arg_group='Network')
     register_cli_argument(scope, 'secrets', multi_ids_type, help='One or many Key Vault secrets as JSON strings or files via \'@<file path>\' containing \'[{ "sourceVault": { "id": "value" }, "vaultCertificates": [{ "certificateUrl": "value", "certificateStore": "cert store name (only on windows)"}] }]\'', type=file_type, completer=FilesCompleter())
-    register_cli_argument(scope, 'os_caching', options_list=['--storage-caching', '--os-caching'], arg_group='Storage', help='Storage caching type for the VM OS disk.', **enum_choice_list(CachingTypes))
-    register_cli_argument(scope, 'data_caching', options_list=['--data-caching'], arg_group='Storage', help='Storage caching type for the VM data disk(s).', **enum_choice_list(CachingTypes))
+    register_cli_argument(scope, 'os_caching', options_list=['--storage-caching', '--os-disk-caching'], arg_group='Storage', help='Storage caching type for the VM OS disk.', **enum_choice_list([CachingTypes.read_only.value, CachingTypes.read_write.value]))
+    register_cli_argument(scope, 'data_caching', options_list=['--data-disk-caching'], arg_group='Storage', help='Storage caching type for the VM data disk(s).', **enum_choice_list(CachingTypes))
 
 register_cli_argument('vm create', 'vm_name', name_arg_type, id_part=None, help='Name of the virtual machine.', validator=process_vm_create_namespace, completer=None)
 register_cli_argument('vm create', 'attach_os_disk', help='Attach an existing OS disk to the VM. Can use the name or ID of a managed disk or the URI to an unmanaged disk VHD.')
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
index 68610f4df87..d0440e807a5 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_template_builder.py
@@ -396,6 +396,7 @@ def _build_data_disks(profile, data_disk_sizes_gb, image_data_disks,
             profile['dataDisks'].append({
                 'lun': image_data_disk.lun,
                 'createOption': "fromImage",
+                'caching': data_caching
             })
         lun = max([d.lun for d in image_data_disks]) + 1 if image_data_disks else 0
         for size in data_disk_sizes_gb:
@@ -540,7 +541,7 @@ def build_vmss_resource(name, naming_prefix, location, tags, overprovision, upgr
             })
         else:
             storage_properties['osDisk']['vhdContainers'] = "[variables('vhdContainers')]"
-    elif storage_profile in [StorageProfile.ManagedPirImage, StorageProfile.ManagedPirImage]:
+    elif storage_profile in [StorageProfile.ManagedPirImage, StorageProfile.ManagedCustomImage]:
         storage_properties['osDisk'] = {
             'createOption': 'FromImage',
             'caching': os_caching,
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
index b39987b4902..98e4a41cba0 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py
@@ -281,15 +281,15 @@ def _validate_vm_create_storage_profile(namespace, for_scale_set=False):
 
     elif namespace.storage_profile == StorageProfile.SAPirImage:
         required = ['image', 'use_unmanaged_disk']
-        forbidden = ['os_type', 'attach_os_disk', 'data_disk_sizes_gb']
+        forbidden = ['os_type', 'data_caching', 'attach_os_disk', 'data_disk_sizes_gb']
 
     elif namespace.storage_profile == StorageProfile.SACustomImage:
         required = ['image', 'os_type', 'use_unmanaged_disk']
-        forbidden = ['attach_os_disk', 'data_disk_sizes_gb']
+        forbidden = ['attach_os_disk', 'data_caching', 'data_disk_sizes_gb']
 
     elif namespace.storage_profile == StorageProfile.SASpecializedOSDisk:
         required = ['os_type', 'attach_os_disk', 'use_unmanaged_disk']
-        forbidden = ['os_disk_name', 'os_caching', 'image', 'storage_account',
+        forbidden = ['os_disk_name', 'os_caching', 'data_caching', 'image', 'storage_account',
                      'storage_container_name', 'data_disk_sizes_gb', 'storage_sku']
 
     else:
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
index 49a7e42d8fb..04e1d09f2a1 100644
--- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
+++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/custom.py
@@ -1682,6 +1682,7 @@ def create_vmss(vmss_name, resource_group_name, image,
     from azure.cli.core._profile import CLOUD
     from azure.mgmt.resource.resources import ResourceManagementClient
     from azure.mgmt.resource.resources.models import DeploymentProperties, TemplateLink
+    from azure.mgmt.compute.models import CachingTypes
 
     network_id_template = resource_id(
         subscription=get_subscription_id(), resource_group=resource_group_name,
@@ -1691,6 +1692,7 @@ def create_vmss(vmss_name, resource_group_name, image,
     tags = tags or {}
     os_disk_name = os_disk_name or 'osdisk_{}'.format(random_string(10))
     storage_container_name = storage_container_name or 'vhds'
+    os_caching = os_caching or CachingTypes.read_write.value
 
     # Build up the ARM template
     master_template = ArmTemplateBuilder()
diff --git a/src/command_modules/azure-cli-vm/tests/test_vm_commands.py b/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
index 6388ece4213..ab5bbe0c38a 100644
--- a/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
+++ b/src/command_modules/azure-cli-vm/tests/test_vm_commands.py
@@ -1443,7 +1443,7 @@ def body(self):
 
         self.cmd('vmss create --image Debian --admin-password Test1234@! -l westus'
                  ' -g {} -n {} --disable-overprovision --instance-count {}'
-                 ' --os-caching {} --upgrade-policy-mode {}'
+                 ' --os-disk-caching {} --upgrade-policy-mode {}'
                  ' --authentication-type password --admin-username myadmin --public-ip-address {}'
                  ' --data-disk-sizes-gb 1 --vm-sku Standard_D2_v2'
                  .format(self.resource_group, vmss_name, instance_count, caching,