diff --git a/perfkitbenchmarker/providers/azure/azure_flexible_server.py b/perfkitbenchmarker/providers/azure/azure_flexible_server.py index 4ddc8ae811..188b94d03a 100644 --- a/perfkitbenchmarker/providers/azure/azure_flexible_server.py +++ b/perfkitbenchmarker/providers/azure/azure_flexible_server.py @@ -122,67 +122,26 @@ def GetResourceMetadata(self) -> dict[str, Any]: def _Create(self) -> None: """Creates the Azure database instance.""" - ha_flag = ENABLE_HA if self.spec.high_availability else DISABLE_HA - # Consider adding --zone and maybe --standby-zone for parity with GCP. - cmd = [ - azure.AZURE_PATH, - self.GetAzCommandForEngine(), - self.SERVER_TYPE, - 'create', - '--resource-group', - self.resource_group.name, - '--name', - self.instance_id, - '--location', - self.region, - '--high-availability', - ha_flag, - '--admin-user', - self.spec.database_username, - '--admin-password', - self.spec.database_password, - '--storage-size', - str(self.spec.db_disk_spec.disk_size), - '--sku-name', - self.spec.db_spec.machine_type, - '--tier', - self.spec.db_tier, - '--version', - self.spec.engine_version, - '--yes', - ] if ( self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES and self.storage_type == azure_disk.PREMIUM_STORAGE_V2 ): - # Postgres PremiumV2_LRS requires creation through REST API. - self._CreatePremiumSsdV2() - return - if self.storage_type: - cmd.extend([ - '--storage-type', - self.storage_type, - ]) - if ( - self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL - and self.spec.db_disk_spec.provisioned_iops is not None - ): - cmd.extend([ - '--auto-scale-iops', - 'Disabled', - ]) - if self.spec.db_disk_spec.provisioned_iops: - cmd.extend([ - '--iops', - str(self.spec.db_disk_spec.provisioned_iops), - ]) - if self.spec.db_disk_spec.provisioned_throughput: - cmd.extend([ - '--throughput', - str(self.spec.db_disk_spec.provisioned_throughput), - ]) - - vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT) + try: + self._CreatePremiumSsdV2() + return + except errors.Resource.CreationError as e: + if self._IsFallbackableError(e): + logging.warning( + 'Failed to create Postgres PremiumV2_LRS instance: %s. ' + 'Falling back to Premium_LRS.', + e, + ) + self._FallbackToPremiumLrs() + # Fall through to _CreateStandard() + else: + raise + + self._CreateStandard() @vm_util.Retry( timeout=CREATE_AZURE_DB_TIMEOUT, @@ -257,18 +216,114 @@ def _CreatePremiumSsdV2(self) -> None: json.dumps(body), '--verbose', ] - _, stderr, _ = vm_util.IssueCommand( + stdout, stderr, retcode = vm_util.IssueCommand( cmd, timeout=CREATE_AZURE_DB_TIMEOUT, raise_on_failure=False ) + if retcode != 0: + raise errors.Resource.CreationError( + f'Azure REST API call failed (retcode={retcode}). ' + f'Stdout: {stdout}, Stderr: {stderr}' + ) match = re.search(r"'Azure-AsyncOperation': '([^']*)'", stderr) if not match: raise errors.Resource.CreationError( 'Could not find Azure-AsyncOperation header in az rest verbose' - ' output.' + f' output. Stdout: {stdout}, Stderr: {stderr}' ) async_operation_url = match.group(1) self._WaitForAzureAsyncOperation(async_operation_url) + def _IsFallbackableError(self, error: Exception) -> bool: + """Returns True if the error is fallbackable.""" + error_str = str(error) + return ( + 'SkuNotAvailable' in error_str + or 'ServerDropping' in error_str + or 'AvailabilityZone' in error_str + or 'not enough resources' in error_str.lower() + ) + + def _FallbackToPremiumLrs(self) -> None: + """Falls back to Premium_LRS storage.""" + logging.info('Initiating deletion of failed instance %s', self.instance_id) + try: + self._Delete() + except Exception as e: # pylint: disable=broad-except + logging.warning( + 'Failed to delete failed instance %s: %s', self.instance_id, e + ) + + old_instance_id = self.instance_id + self.instance_id = f'{self.instance_id}-fb' + self.storage_type = azure_disk.PREMIUM_STORAGE + logging.info( + 'Falling back from %s (PremiumV2_LRS) to %s (Premium_LRS)', + old_instance_id, + self.instance_id, + ) + + def _CreateStandard(self) -> None: + """Creates standard Azure database instance (MySQL or Postgres Premium).""" + ha_flag = ENABLE_HA if self.spec.high_availability else DISABLE_HA + cmd = [ + azure.AZURE_PATH, + self.GetAzCommandForEngine(), + self.SERVER_TYPE, + 'create', + '--resource-group', + self.resource_group.name, + '--name', + self.instance_id, + '--location', + self.region, + '--high-availability', + ha_flag, + '--admin-user', + self.spec.database_username, + '--admin-password', + self.spec.database_password, + '--storage-size', + str(self.spec.db_disk_spec.disk_size), + '--sku-name', + self.spec.db_spec.machine_type, + '--tier', + self.spec.db_tier, + '--version', + self.spec.engine_version, + '--yes', + ] + if self.storage_type: + cmd.extend([ + '--storage-type', + self.storage_type, + ]) + if ( + self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL + and self.spec.db_disk_spec.provisioned_iops is not None + ): + cmd.extend([ + '--auto-scale-iops', + 'Disabled', + ]) + + is_postgres = self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES + is_premium_v2 = self.storage_type == azure_disk.PREMIUM_STORAGE_V2 + + if self.spec.db_disk_spec.provisioned_iops: + if not is_postgres or is_premium_v2: + cmd.extend([ + '--iops', + str(self.spec.db_disk_spec.provisioned_iops), + ]) + if self.spec.db_disk_spec.provisioned_throughput: + if not is_postgres or is_premium_v2: + cmd.extend([ + '--throughput', + str(self.spec.db_disk_spec.provisioned_throughput), + ]) + + vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT) + def GetAzCommandForEngine(self) -> str: engine = self.spec.engine if engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES: diff --git a/tests/providers/azure/azure_flexible_server_test.py b/tests/providers/azure/azure_flexible_server_test.py index e4c8800a50..058547ebc8 100644 --- a/tests/providers/azure/azure_flexible_server_test.py +++ b/tests/providers/azure/azure_flexible_server_test.py @@ -196,7 +196,7 @@ def testCreatePostgres(self): with self.subTest(name='StorageType'): self.assertIn('--storage-type Premium_LRS', cmd) with self.subTest(name='Iops'): - self.assertIn('--iops 1000', cmd) + self.assertNotIn('--iops', cmd) with self.subTest(name='Metadata'): self.assertEqual(metadata['disk_type'], 'Premium_LRS') @@ -415,6 +415,78 @@ def testCreatePostgresPremiumV2NoAsyncHeader(self): with self.assertRaises(errors.Resource.CreationError): bm_spec.relational_db._Create() + def testCreatePostgresPremiumV2Fallback(self): + mock_cmd = self.MockIssueCommand({ + 'rest --method PUT': [( + '', + "'Azure-AsyncOperation': 'https://management.azure.com/async_op'", + 0, + )], + 'rest --method GET': [( + '{"status": "Failed", "error": {"code": "ServerDropping"}}', + '', + 0, + )], + 'postgres flexible-server delete': [('', '', 0)], + 'postgres flexible-server create': [('', '', 0)], + }) + yaml_spec = inspect.cleandoc(f""" + sysbench: + relational_db: + cloud: {provider_info.AZURE} + engine: {sql_engine_utils.FLEXIBLE_SERVER_POSTGRES} + engine_version: '13' + db_tier: GeneralPurpose + db_spec: + {provider_info.AZURE}: + machine_type: Standard_D2ds_v4 + zone: eastus + db_disk_spec: + {provider_info.AZURE}: + disk_size: 128 + disk_type: PremiumV2_LRS + provisioned_iops: 1000 + provisioned_throughput: 200 + vm_groups: + clients: + vm_spec: *default_dual_core + disk_spec: *default_500_gb + """) + bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml( + yaml_spec, 'sysbench' + ) + bm_spec.ConstructRelationalDb() + db = bm_spec.relational_db + db.resource_group.name = 'az_resource' + + db._Create() + + self.assertEqual(mock_cmd.func_to_mock.call_count, 4) + put_call, get_call, delete_call, create_call = ( + mock_cmd.func_to_mock.call_args_list + ) + + put_cmd = ' '.join(put_call[0][0]) + self.assertIn('rest --method PUT', put_cmd) + self.assertIn('PremiumV2_LRS', put_cmd) + + get_cmd = ' '.join(get_call[0][0]) + self.assertIn('rest --method GET', get_cmd) + + delete_cmd = ' '.join(delete_call[0][0]) + self.assertIn('postgres flexible-server delete', delete_cmd) + self.assertIn('--name pkb-db-instance-test_run_uri', delete_cmd) + + create_cmd = ' '.join(create_call[0][0]) + self.assertIn('postgres flexible-server create', create_cmd) + self.assertIn('--name pkb-db-instance-test_run_uri-fb', create_cmd) + self.assertIn('--storage-type Premium_LRS', create_cmd) + self.assertNotIn('--iops', create_cmd) + self.assertNotIn('--throughput', create_cmd) + + self.assertEqual(db.instance_id, 'pkb-db-instance-test_run_uri-fb') + self.assertEqual(db.storage_type, 'Premium_LRS') + if __name__ == '__main__': unittest.main()