Add regional support for GCE MIGs

PiperOrigin-RevId: 934454503

Author: pmkc

perfkitbenchmarker/managed_vm_group.py

@@ -59,13 +59,19 @@ def __init__(
     self._vms: dict[str, virtual_machine.BaseVirtualMachine] = {}
     self._deleted_vms: list[virtual_machine.BaseVirtualMachine] = []
 
+    self.region: str | None = None
+    self.zone: str | None = None
+
     self.last_operation_start_time: float | None = None
     self.last_ready_time: float | None = None
 
   def GetResourceMetadata(self) -> dict[Any, Any]:
     metadata = super().GetResourceMetadata().copy()
     vm = (list(self.vms) + self._deleted_vms + [self.vm_config])[0]
     metadata.update(vm.GetResourceMetadata())
+    metadata['region'] = self.region
+    metadata['zone'] = self.zone
+    metadata['is_regional'] = self.zone is None
     return metadata
 
   @property

perfkitbenchmarker/providers/gcp/gce_managed_instance_group.py

@@ -14,6 +14,7 @@
 """GCE Managed Instance Group resource."""
 
 import json
+import re
 from typing import Any, cast
 
 from perfkitbenchmarker import errors
@@ -101,6 +102,7 @@ class GceManagedInstanceGroup(managed_vm_group.BaseManagedVmGroup):
   """GCE Managed Instance Group."""
 
   CLOUD = provider_info.GCP
+  instance_template: GceInstanceTemplate
 
   def __init__(
       self,
@@ -112,10 +114,18 @@ def __init__(
         gce_virtual_machine.GceVirtualMachine, self.vm_config
     )
     self.project = self.vm_config.project
-    # TODO(pclay): Add support for regional managed instance groups.
-    # It's unclear how multiple zones would be plumbed through BaseVmSpec.
-    self.zone = self.vm_config.zone
-    self.region = util.GetRegionFromZone(self.zone)
+    # in theory we could use the users defaults, but this is cleaner.
+    assert self.vm_config.zone
+    if util.IsZone(self.vm_config.zone):
+      self.zone = self.vm_config.zone
+      self.region = util.GetRegionFromZone(self.zone)
+    elif util.IsRegion(self.vm_config.zone):
+      self.region = self.vm_config.zone
+      self.zone = None
+    else:
+      raise ValueError(
+          f'Unsupported zone: {self.vm_config.zone}. Must be a zone or region.'
+      )
 
     self.instance_template = GceInstanceTemplate(
         self.vm_config, self.region, name=self.name
@@ -126,9 +136,13 @@ def _CreateDependencies(self):
     self.instance_template.Create()
 
   def _GcloudCmd(self, *args) -> util.GcloudCommand:
-    return util.GcloudCommand(
+    cmd = util.GcloudCommand(
         self, 'compute', 'instance-groups', 'managed', *args
     )
+    # GcloudCommand does not have support for regional resources.
+    if not self.zone:
+      cmd.flags['region'] = self.region
+    return cmd
 
   def _Create(self):
     cmd = self._GcloudCmd(
@@ -176,7 +190,14 @@ def _GetCurrentVms(self) -> list[VmReference]:
           f'{stderr}\n'
       )
     instances = json.loads(stdout)
-    return [VmReference(name=instance['name']) for instance in instances]
+    references = []
+    for instance in instances:
+      match = re.search(
+          r'/zones/([^/]+)/instances/([^/]+)', instance['instance']
+      )
+      assert match
+      references.append(VmReference(name=match.group(2), zone=match.group(1)))
+    return references
 
   def _AddVms(self, num_vms_to_add: int):
     cmd = self._GcloudCmd('create-instance', self.name)

tests/providers/gcp/gce_managed_instance_group_test.py

@@ -3,6 +3,7 @@
 from unittest import mock
 
 from absl import flags
+from perfkitbenchmarker import virtual_machine
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.configs import vm_group_decoders
 from perfkitbenchmarker.providers.gcp import gce_managed_instance_group
@@ -18,6 +19,7 @@ class GceManagedInstanceGroupTest(pkb_common_test_case.PkbCommonTestCase):
   def setUp(self):
     super().setUp()
     FLAGS.run_uri = 'test_run'
+    virtual_machine.BaseVirtualMachine._instance_counter = 0
     self.mock_cmd = self.MockIssueCommand(
         {
             '': [('', '', 0)],
@@ -27,13 +29,13 @@ def setUp(self):
   @mock.patch.object(gcp_utils, 'GetRegionFromZone', return_value='us-central1')
   @mock.patch.object(gce_virtual_machine.gce_network.GceFirewall, 'GetFirewall')
   @mock.patch.object(gce_virtual_machine.gce_network.GceNetwork, 'GetNetwork')
-  def TestMig(self, mock_get_network, *_):
+  def TestMig(self, mock_get_network, *_, zone='us-central1-c'):
     mock_get_network.return_value.placement_group.name = 'test_placement_group'
     vm_config = pkb_common_test_case.TestGceVirtualMachine(
         gce_virtual_machine.GceVmSpec(
             'test_component',
             machine_type='n1-standard-4',
-            zone='us-central1-c',
+            zone=zone,
         )
     )
     return gce_managed_instance_group.GceManagedInstanceGroup(
@@ -47,7 +49,7 @@ def TestMig(self, mock_get_network, *_):
     )
 
   def testCreate(self, *_):
-    mig = self.TestMig()
+    mig = self.TestMig(zone='us-central1-c')
     mig._Create()
     self.assertIn(
         'gcloud compute instance-groups managed create pkb-test_run-0'
@@ -58,6 +60,18 @@ def testCreate(self, *_):
         self.mock_cmd.all_commands,
     )
 
+  def testCreateRegional(self, *_):
+    mig = self.TestMig(zone='us-central1')
+    mig._Create()
+    self.assertIn(
+        'gcloud compute instance-groups managed create pkb-test_run-0'
+        ' --template'
+        ' projects/test_project/regions/us-central1/instanceTemplates/pkb-test_run-0'
+        ' --size 1 --format json --project test_project --quiet --region'
+        ' us-central1',
+        self.mock_cmd.all_commands,
+    )
+
   # SSH keys
   @mock.patch.object(builtins, 'open')
   @mock.patch.object(vm_util, 'NamedTemporaryFile')
@@ -68,7 +82,7 @@ def testCreateDependencies(self, mock_named_temporary_file, _):
     )
     mig._CreateDependencies()
     self.assertIn(
-        'gcloud compute instance-templates create pkb-test_run-1'
+        'gcloud compute instance-templates create pkb-test_run-0'
         ' --instance-template-region us-central1 --format json --labels '
         ' --machine-type n1-standard-4 --maintenance-policy TERMINATE'
         ' --metadata enable-oslogin=FALSE,vm_nature=ephemeral'
@@ -78,6 +92,25 @@ def testCreateDependencies(self, mock_named_temporary_file, _):
         self.mock_cmd.all_commands,
     )
 
+  def testGetCurrentVms(self):
+    list_instances_response = (
+        '[{"name": "pkb-test_run-0-0", "instance":'
+        ' "https://www.googleapis.com/compute/v1/projects/test_project/'
+        'zones/us-central1-c/instances/pkb-test_run-0-0"}]'
+    )
+    self.MockIssueCommand(
+        {
+            'gcloud compute instance-groups managed list-instances': [
+                (list_instances_response, '', 0)
+            ]
+        }
+    )
+    mig = self.TestMig()
+    vms = mig._GetCurrentVms()
+    self.assertEqual(len(vms), 1)
+    self.assertEqual(vms[0].name, 'pkb-test_run-0-0')
+    self.assertEqual(vms[0].zone, 'us-central1-c')
+
 
 if __name__ == '__main__':
   unittest.main()