From 3b633f060eaa78b6a2122753432c28d6aef4711e Mon Sep 17 00:00:00 2001 From: Raymond Lam Date: Fri, 5 Jan 2024 12:06:19 -0800 Subject: [PATCH] Support setting up temp db for Azure. Azure uses D drive for temp disk. PiperOrigin-RevId: 596052324 --- perfkitbenchmarker/disk_strategies.py | 2 +- .../providers/azure/azure_disk_strategies.py | 43 +++++++++++++++++++ .../providers/azure/azure_relational_db.py | 4 +- .../providers/azure/azure_virtual_machine.py | 5 +-- .../sqlserver_iaas_relational_db.py | 7 ++- 5 files changed, 50 insertions(+), 11 deletions(-) create mode 100644 perfkitbenchmarker/providers/azure/azure_disk_strategies.py diff --git a/perfkitbenchmarker/disk_strategies.py b/perfkitbenchmarker/disk_strategies.py index 1d8fb069ee..e2912a2242 100644 --- a/perfkitbenchmarker/disk_strategies.py +++ b/perfkitbenchmarker/disk_strategies.py @@ -230,7 +230,7 @@ def PrepareLinuxScratchDisk( scratch_disk: Union[disk.BaseDisk, disk.StripedDisk], disk_spec: disk.BaseDiskSpec, ) -> None: - """Helper method to format and mount scratch disk. + """Format and mount scratch disk. Args: vm: Linux Virtual Machine to create scratch disk on. diff --git a/perfkitbenchmarker/providers/azure/azure_disk_strategies.py b/perfkitbenchmarker/providers/azure/azure_disk_strategies.py new file mode 100644 index 0000000000..87522fed3b --- /dev/null +++ b/perfkitbenchmarker/providers/azure/azure_disk_strategies.py @@ -0,0 +1,43 @@ +# Copyright 2024 PerfKitBenchmarker Authors. All rights reserved. +# +# Licensed 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. +"""Module containing strategies to prepare disks. + +This module abstract out the disk algorithm for formatting and creating +scratch disks. +""" + +from typing import Any +from absl import flags +from perfkitbenchmarker import disk_strategies +from perfkitbenchmarker.providers.azure import azure_disk + +FLAGS = flags.FLAGS + + +virtual_machine = Any # pylint: disable=invalid-name + + +class AzurePrepareScratchDiskStrategy( + disk_strategies.PrepareScratchDiskStrategy +): + """Strategies to prepare scratch disk on AWS.""" + + def GetLocalSSDNames(self) -> list[str]: + # This is only for Ls-Series local ssd + # Temp disk is automatically setup on Azure and default to disk D. + return ['Microsoft NVMe Direct Disk'] + + def PrepareTempDbDisk(self, vm: 'virtual_machine.BaseVirtualMachine'): + if azure_disk.HasTempDrive(vm.machine_type): + vm.RemoteCommand('mkdir D:\\TEMPDB') diff --git a/perfkitbenchmarker/providers/azure/azure_relational_db.py b/perfkitbenchmarker/providers/azure/azure_relational_db.py index 4760721dbd..6695bf2e78 100644 --- a/perfkitbenchmarker/providers/azure/azure_relational_db.py +++ b/perfkitbenchmarker/providers/azure/azure_relational_db.py @@ -62,9 +62,7 @@ class AzureSQLServerIAASRelationalDb( """A AWS IAAS database resource.""" CLOUD = provider_info.AZURE - - def MoveSQLServerTempDb(self): - """Moves the SQL Server temporary database to LocalSSD.""" + TEMPDB_DISK_LETTER = 'D' class AzurePostgresIAASRelationalDb( diff --git a/perfkitbenchmarker/providers/azure/azure_virtual_machine.py b/perfkitbenchmarker/providers/azure/azure_virtual_machine.py index c82f44522e..d1877894e0 100644 --- a/perfkitbenchmarker/providers/azure/azure_virtual_machine.py +++ b/perfkitbenchmarker/providers/azure/azure_virtual_machine.py @@ -24,7 +24,6 @@ operate on the VM: boot, shutdown, etc. """ - import abc import collections import itertools @@ -39,7 +38,6 @@ from absl import flags from perfkitbenchmarker import custom_virtual_machine_spec from perfkitbenchmarker import disk -from perfkitbenchmarker import disk_strategies from perfkitbenchmarker import errors from perfkitbenchmarker import linux_virtual_machine from perfkitbenchmarker import os_types @@ -52,6 +50,7 @@ from perfkitbenchmarker.configs import option_decoders from perfkitbenchmarker.providers import azure from perfkitbenchmarker.providers.azure import azure_disk +from perfkitbenchmarker.providers.azure import azure_disk_strategies from perfkitbenchmarker.providers.azure import azure_network from perfkitbenchmarker.providers.azure import util from six.moves import range @@ -1056,7 +1055,7 @@ def CreateScratchDisk(self, _, disk_spec): disks.append(data_disk) scratch_disk = self._CreateScratchDiskFromDisks(disk_spec, disks) - disk_strategies.PrepareScratchDiskStrategy().PrepareScratchDisk( + azure_disk_strategies.AzurePrepareScratchDiskStrategy().PrepareScratchDisk( self, scratch_disk, disk_spec ) diff --git a/perfkitbenchmarker/sqlserver_iaas_relational_db.py b/perfkitbenchmarker/sqlserver_iaas_relational_db.py index 0a3c375805..d47d579d21 100644 --- a/perfkitbenchmarker/sqlserver_iaas_relational_db.py +++ b/perfkitbenchmarker/sqlserver_iaas_relational_db.py @@ -47,11 +47,10 @@ DEFAULT_ENGINE_VERSION = sql_engine_utils.SQLSERVER_ENTERPRISE -TEMPDB_DISK_LETTER = "T" - class SQLServerIAASRelationalDb(iaas_relational_db.IAASRelationalDb): """Object representing a IAAS relational database Service.""" + TEMPDB_DISK_LETTER = "T" ENGINE = sql_engine_utils.SQLSERVER @@ -123,7 +122,7 @@ def MoveSQLServerTempDb(self): if drive ] - if TEMPDB_DISK_LETTER in drive_list: + if self.TEMPDB_DISK_LETTER in drive_list: stdout, _ = vm.RemoteCommand( 'sqlcmd -h -1 -Q "SET NOCOUNT ' " ON; select f.name + CASE WHEN " @@ -145,7 +144,7 @@ def MoveSQLServerTempDb(self): 'sqlcmd -Q "ALTER DATABASE tempdb ' "MODIFY FILE (NAME = [{}], " "FILENAME = '{}:\\TEMPDB\\{}');\"".format( - tmp_db_name, TEMPDB_DISK_LETTER, tmp_db_file + tmp_db_name, self.TEMPDB_DISK_LETTER, tmp_db_file ) )