Add Azure Flexible Server metrics implementation.

bvliu · copybara-github · commit 0e90a189efca · 2025-12-19T07:57:53.000-08:00
PiperOrigin-RevId: 846734952
diff --git a/perfkitbenchmarker/providers/azure/azure_flexible_server.py b/perfkitbenchmarker/providers/azure/azure_flexible_server.py
@@ -20,6 +20,7 @@
 """
 
 import datetime
+import json
 import logging
 import time
 from typing import Any, Tuple
@@ -28,11 +29,13 @@
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import provider_info
 from perfkitbenchmarker import relational_db
+from perfkitbenchmarker import sample
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.providers import azure
 from perfkitbenchmarker.providers.azure import azure_disk
 from perfkitbenchmarker.providers.azure import azure_relational_db
+from perfkitbenchmarker.providers.azure import util
 
 DEFAULT_DATABASE_NAME = 'database'
 
@@ -70,6 +73,9 @@ class AzureFlexibleServer(azure_relational_db.AzureRelationalDb):
       sql_engine_utils.FLEXIBLE_SERVER_POSTGRES,
       sql_engine_utils.FLEXIBLE_SERVER_MYSQL,
   ]
+  # Metrics are processed in 5 minute batches according to
+  # https://learn.microsoft.com/en-us/azure/postgresql/flexible-server/concepts-monitoring.
+  METRICS_COLLECTION_DELAY_SECONDS = 300
 
   def __init__(self, relational_db_spec: Any):
     super().__init__(relational_db_spec)
@@ -277,3 +283,122 @@ def _ApplyDbFlags(self) -> None:
         )
 
     self._Reboot()
+
+  def _GetResourceProvider(self) -> str:
+    if self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL:
+      return 'Microsoft.DBforMySQL'
+    elif self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES:
+      return 'Microsoft.DBforPostgreSQL'
+    else:
+      raise NotImplementedError(f'Unsupported engine {self.spec.engine}')
+
+  def _GetResourceId(self) -> str:
+    return (
+        f'/subscriptions/{util.GetSubscriptionId()}/resourceGroups/'
+        f'{self.resource_group.name}/providers/'
+        f'{self._GetResourceProvider()}/flexibleServers/{self.instance_id}'
+    )
+
+  def _GetMetricsToCollect(self) -> list[relational_db.MetricSpec]:
+    """Returns a list of metrics to collect."""
+    # pyformat: disable
+    if self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL:
+      return [
+          relational_db.MetricSpec('cpu_percent', 'cpu_utilization', '%', None),
+          relational_db.MetricSpec('io_consumption_percent', 'io_consumption_percent', '%', None),
+          relational_db.MetricSpec('storage_io_count', 'storage_io_count', 'iops', None),
+          relational_db.MetricSpec('storage_used', 'disk_bytes_used', 'GB', lambda x: x / (1024 * 1024 * 1024)),
+      ]
+    else:
+      return [
+          relational_db.MetricSpec('cpu_percent', 'cpu_utilization', '%', None),
+          relational_db.MetricSpec('read_iops', 'disk_read_iops', 'iops', None),
+          relational_db.MetricSpec('write_iops', 'disk_write_iops', 'iops', None),
+          relational_db.MetricSpec('read_throughput', 'disk_read_throughput', 'MB/s', lambda x: x / (1024 * 1024)),
+          relational_db.MetricSpec('write_throughput', 'disk_write_throughput', 'MB/s', lambda x: x / (1024 * 1024)),
+          relational_db.MetricSpec('storage_used', 'disk_bytes_used', 'GB', lambda x: x / (1024 * 1024 * 1024)),
+      ]
+    # pyformat: enable
+
+  @vm_util.Retry(poll_interval=60, max_retries=5, retryable_exceptions=KeyError)
+  def _CollectProviderMetric(
+      self,
+      metric: relational_db.MetricSpec,
+      start_time: datetime.datetime,
+      end_time: datetime.datetime,
+      collect_percentiles: bool = False,
+  ) -> list[sample.Sample]:
+    """Collects metrics from Azure Monitor."""
+    if end_time - start_time < datetime.timedelta(minutes=1):
+      logging.warning(
+          'Not collecting metrics since end time %s is within 1 minute of start'
+          ' time %s.',
+          end_time,
+          start_time,
+      )
+      return []
+    metric_name = metric.provider_name
+    logging.info(
+        'Collecting metric %s for instance %s', metric_name, self.instance_id
+    )
+    cmd = [
+        azure.AZURE_PATH,
+        'monitor',
+        'metrics',
+        'list',
+        '--resource',
+        self._GetResourceId(),
+        '--metric',
+        metric_name,
+        '--start-time',
+        start_time.astimezone(datetime.timezone.utc).strftime(
+            relational_db.METRICS_TIME_FORMAT
+        ),
+        '--end-time',
+        end_time.astimezone(datetime.timezone.utc).strftime(
+            relational_db.METRICS_TIME_FORMAT
+        ),
+        '--interval',
+        'pt1m',
+        '--aggregation',
+        'Average',
+    ]
+    try:
+      stdout, _ = vm_util.IssueRetryableCommand(cmd)
+    except errors.VmUtil.IssueCommandError as e:
+      logging.warning(
+          'Could not collect metric %s for instance %s: %s',
+          metric.provider_name,
+          self.instance_id,
+          e,
+      )
+      return []
+    response = json.loads(stdout)
+    if (
+        not response
+        or not response['value']
+        or not response['value'][0]['timeseries']
+    ):
+      logging.warning('No timeseries for metric %s', metric_name)
+      return []
+
+    datapoints = response['value'][0]['timeseries'][0]['data']
+    if not datapoints:
+      logging.warning('No datapoints for metric %s', metric_name)
+      return []
+
+    points = []
+    for dp in datapoints:
+      if dp['average'] is None:
+        continue
+      value = dp['average']
+      if metric.conversion_func:
+        value = metric.conversion_func(value)
+      points.append((
+          datetime.datetime.fromisoformat(dp['timeStamp']),
+          value,
+      ))
+
+    return self._CreateSamples(
+        points, metric.sample_name, metric.unit, collect_percentiles
+    )
diff --git a/perfkitbenchmarker/providers/azure/util.py b/perfkitbenchmarker/providers/azure/util.py
@@ -280,3 +280,11 @@ def GetMachineSeriesNumber(machine_type: str) -> int:
     return int(match.group(1))
   # Azure only adds a v after the first series,
   return 1
+
+
+def GetSubscriptionId():
+  """Gets the current Azure subscription ID."""
+  stdout, _ = vm_util.IssueRetryableCommand(
+      [AZURE_PATH, 'account', 'show', '--query', 'id', '-o', 'tsv']
+  )
+  return stdout.strip()
diff --git a/tests/providers/azure/azure_flexible_server_test.py b/tests/providers/azure/azure_flexible_server_test.py
@@ -1,4 +1,6 @@
+import datetime
 import inspect
+import json
 import unittest
 
 from absl import flags
@@ -8,11 +10,109 @@
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.providers.azure import azure_network
+from perfkitbenchmarker.providers.azure import util
 from tests import pkb_common_test_case
 
+
 FLAGS = flags.FLAGS
 
 
+class AzureFlexibleServerMetricsTest(pkb_common_test_case.PkbCommonTestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.enter_context(
+        mock.patch.object(vm_util, 'IssueCommand', return_value=('', '', ''))
+    )
+    FLAGS.run_uri = '123'
+    FLAGS.cloud = provider_info.AZURE
+    FLAGS['db_engine'].parse(sql_engine_utils.FLEXIBLE_SERVER_POSTGRES)
+    test_spec = inspect.cleandoc("""
+    sysbench:
+      relational_db:
+        engine: postgres
+        engine_version: '13'
+        database_username: user
+        database_password: password
+        high_availability: False
+        db_spec:
+          Azure:
+            machine_type: Standard_D2s_v3
+            zone: westus2
+        db_disk_spec:
+          Azure:
+            disk_size: 128
+        vm_groups:
+          clients:
+            vm_spec:
+              Azure:
+                machine_type: Standard_B4ms
+                zone: westus2
+    """)
+    self.spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
+        yaml_string=test_spec, benchmark_name='sysbench'
+    )
+    self.spec.ConstructRelationalDb()
+    self.server = self.spec.relational_db
+    self.server.resource_group = mock.Mock()
+    self.server.resource_group.name = 'test-group'
+    self.enter_context(
+        mock.patch.object(util, 'GetSubscriptionId', return_value='test-sub')
+    )
+
+  def testCollectMetrics(self):
+    # Mock the response from Azure Monitor
+    mock_response = {
+        'value': [{
+            'timeseries': [{
+                'data': [
+                    {
+                        'timeStamp': '2025-11-26T10:00:00Z',
+                        'average': 10.0,
+                    },
+                    {
+                        'timeStamp': '2025-11-26T10:01:00Z',
+                        'average': 20.0,
+                    },
+                ]
+            }]
+        }]
+    }
+    self.enter_context(
+        mock.patch.object(
+            vm_util,
+            'IssueRetryableCommand',
+            return_value=(json.dumps(mock_response), ''),
+        )
+    )
+
+    start_time = datetime.datetime(2025, 11, 26, 10, 0, 0)
+    end_time = datetime.datetime(2025, 11, 26, 10, 1, 0)
+    samples = self.server.CollectMetrics(start_time, end_time)
+
+    # Check the number of samples returned (4 per metric * 6 metrics)
+    self.assertLen(samples, 24)
+
+    # Spot check a few sample values
+    sample_names = [s.metric for s in samples]
+    self.assertIn('cpu_utilization_average', sample_names)
+    self.assertIn('cpu_utilization_min', sample_names)
+    self.assertIn('cpu_utilization_max', sample_names)
+    self.assertIn('disk_read_iops_average', sample_names)
+
+    cpu_avg = next(s for s in samples if s.metric == 'cpu_utilization_average')
+    self.assertEqual(cpu_avg.value, 15.0)
+    self.assertEqual(cpu_avg.unit, '%')
+
+    cpu_min = next(s for s in samples if s.metric == 'cpu_utilization_min')
+    self.assertEqual(cpu_min.value, 10.0)
+    self.assertEqual(cpu_min.unit, '%')
+
+    cpu_max = next(s for s in samples if s.metric == 'cpu_utilization_max')
+    self.assertEqual(cpu_max.value, 20.0)
+    self.assertEqual(cpu_max.unit, '%')
+
+
 class AzureFlexibleServerCreateTestCase(pkb_common_test_case.PkbCommonTestCase):
 
   def setUp(self):
