In addition to regional VM groups support listing all zones of group explicitly.

This is necessary on AWS when using manual zonal subnets to know where to create them.

PiperOrigin-RevId: 934477697

Author: pmkc

perfkitbenchmarker/benchmark_spec.py

@@ -823,13 +823,17 @@ def ConstructVirtualMachineGroup(
     # If the VM group is managed, create a managed VM group and return.
     if is_managed:
       managed_vm_group_class = managed_vm_group.GetManagedVmGroupClass(cloud)
-      # TODO(pclay): support multiple zones:
-      if FLAGS.zone:
-        assert len(FLAGS.zone) == 1, 'Managed VM groups only support one zone.'
-        group_spec.vm_spec.zone = FLAGS.zone[0]
-      vm_config = self._CreateVirtualMachine(group_spec.vm_spec, os_type, cloud)
+      zones = FLAGS.zone or [group_spec.vm_spec.zone]
+      vm_configs = []
+      # VM groups needs a VM for each zone to create subnets in each zone.
+      for zone in zones:
+        spec = copy.copy(group_spec.vm_spec)
+        spec.zone = zone
+        vm_config = self._CreateVirtualMachine(spec, os_type, cloud)
+        vm_config.zone = zone
+        vm_configs.append(vm_config)
       group = managed_vm_group_class(
-          group_spec, vm_config
+          group_spec, vm_configs
       )  # pytype: disable=not-instantiable
       self.managed_vm_groups[group_name] = group
       # Report resource provisioning times.

perfkitbenchmarker/managed_vm_group.py

@@ -32,21 +32,50 @@ def GetManagedVmGroupClass(cloud):
 
 
 class BaseManagedVmGroup(resource.BaseResource):
-  """Base class representing a Managed VM group."""
+  """Base class representing a Managed VM group.
+
+  Attributes:
+    name: The name of the managed VM group.
+    region: The region where the managed VM group is located.
+    vms: A sequence of all VMs in the managed VM group.
+
+  Internal Attributes:
+    spec: The spec of the managed VM group.
+    vm_config: A virtual machine instance representing the prototype/template VM
+      configuration used to clone new VMs in the group. Individual VM operations
+      like `_Create()` and `_Delete()` are NOT called on this object (creation
+      and deletion are managed by the VM group at the cloud level), but it
+      defines the template settings.
+    zoned_vm_configs: A list of all template VM instances, one for each
+      configured zone in the group. Used to pre-create and cleanup zone-specific
+      VM dependencies (like firewalls, networks, subnets) in every candidate
+      zone.
+    _vms: A mapping of VM names to active VM objects representing actual
+      instances running in the cloud.
+  """
 
   name: str
+  region: str
 
   RESOURCE_TYPE = 'BaseManagedVmGroup'
   CLOUD = None
 
   def __init__(
       self,
       spec: vm_group_decoders.VmGroupSpec,
-      vm_config: virtual_machine.BaseVirtualMachine,
+      vm_configs: list[virtual_machine.BaseVirtualMachine],
   ):
     super().__init__()
     self.spec: vm_group_decoders.VmGroupSpec = spec
-    self.vm_config = vm_config
+    assert vm_configs, 'No VM configs provided.'
+    # `vm_config` is the primary prototype/template VM configuration.
+    # We do NOT run `_Create` or `_Delete` on it. Instead, we use it to
+    # copy-clone the target VMs.
+    self.vm_config = vm_configs[0]
+    # `zoned_vm_configs` contains a template VM configuration for each zone.
+    # Used to create/delete dependencies (e.g. networks, firewalls) in all
+    # zones.
+    self.zoned_vm_configs = vm_configs
     self.vm_config.metadata['in_managed_vm_group'] = True
     # When we clone the VM config and rename it, our assumptions about the
     # disk names are wrong.
@@ -56,22 +85,37 @@ def __init__(
 
     self.name = self.vm_config.name
     self.vm_count = self.spec.vm_count
+    # `_vms` mapping VM names to active VM objects representing actual instances
+    # running in the cloud. They are populated by copy-cloning `vm_config` and
+    # initializing names/zones. Then passed back to the benchmark spec for
+    # benchmarking.
     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.zones: list[str] = [vm.zone for vm in vm_configs]
 
     self.last_operation_start_time: float | None = None
     self.last_ready_time: float | None = None
 
+  @property
+  def is_regional(self) -> bool:
+    assert self.zones
+    return len(self.zones) > 1 or self.zones[0] == self.region
+
+  @property
+  def zone(self) -> str | None:
+    assert self.zones
+    if self.is_regional:
+      return None
+    return self.zones[0]
+
   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
+    metadata['zones'] = self.zones
+    metadata['is_regional'] = self.is_regional
     return metadata
 
   @property
@@ -80,10 +124,16 @@ def vms(self) -> Sequence[virtual_machine.BaseVirtualMachine]:
     return list(self._vms.values())
 
   def _CreateDependencies(self):
-    self.vm_config._CreateDependencies()  # pylint: disable=protected-access
+    background_tasks.RunThreaded(
+        lambda vm: vm._CreateDependencies(),  # pylint: disable=protected-access
+        self.zoned_vm_configs,
+    )
 
   def _DeleteDependencies(self):
-    self.vm_config._DeleteDependencies()  # pylint: disable=protected-access
+    background_tasks.RunThreaded(
+        lambda vm: vm._DeleteDependencies(),  # pylint: disable=protected-access
+        self.zoned_vm_configs,
+    )
 
   @dataclasses.dataclass
   class VmReference:

perfkitbenchmarker/providers/aws/aws_auto_scaling_group.py

@@ -116,15 +116,17 @@ class AwsAutoScalingGroup(managed_vm_group.BaseManagedVmGroup):
   def __init__(
       self,
       spec: vm_group_decoders.VmGroupSpec,
-      vm_config: virtual_machine.BaseVirtualMachine,
+      vm_configs: list[virtual_machine.BaseVirtualMachine],
   ):
-    super().__init__(spec, vm_config)
+    super().__init__(spec, vm_configs)
     self.vm_config: aws_virtual_machine.AwsVirtualMachine = cast(
         aws_virtual_machine.AwsVirtualMachine, self.vm_config
     )
+    self.zoned_vm_configs: list[aws_virtual_machine.AwsVirtualMachine] = cast(
+        list[aws_virtual_machine.AwsVirtualMachine], self.zoned_vm_configs
+    )
     self.name = self.vm_config.name
     self.region = self.vm_config.region
-    self.zone = self.vm_config.zone
 
     self.launch_template = AwsLaunchTemplate(self.vm_config, name=self.name)
     self.base_cmd = util.AWS_PREFIX + ['autoscaling', '--region', self.region]
@@ -134,7 +136,9 @@ def _CreateDependencies(self):
     self.launch_template.Create()
 
   def _Create(self):
-    subnets = self.vm_config.network.subnet.id
+    subnets = []
+    for vm in self.zoned_vm_configs:
+      subnets.append(vm.network.subnet.id)
     cmd = self.base_cmd + [
         'create-auto-scaling-group',
         '--auto-scaling-group-name',
@@ -150,7 +154,7 @@ def _Create(self):
         '--desired-capacity',
         str(self.vm_count),
         '--vpc-zone-identifier',
-        subnets,
+        ','.join(subnets),
     ]
     if self.vm_config.aws_tags:
       cmd.append('--tags')

perfkitbenchmarker/providers/azure/azure_vm_scaling_set.py

@@ -37,12 +37,18 @@ class AzureVmScalingSet(managed_vm_group.BaseManagedVmGroup):
   def __init__(
       self,
       spec: vm_group_decoders.VmGroupSpec,
-      vm_config: virtual_machine.BaseVirtualMachine,
+      vm_configs: list[virtual_machine.BaseVirtualMachine],
   ):
-    super().__init__(spec, vm_config)
+    super().__init__(spec, vm_configs)
     self.vm_config: azure_virtual_machine.AzureVirtualMachine = cast(
         azure_virtual_machine.AzureVirtualMachine, self.vm_config
     )
+    self.zoned_vm_configs: list[azure_virtual_machine.AzureVirtualMachine] = (
+        cast(
+            list[azure_virtual_machine.AzureVirtualMachine],
+            self.zoned_vm_configs,
+        )
+    )
     # VMSS cannot use preconstructed NICs and IPs.
     # AFAICT there is no reason to pre-create them for VMs either.
     for resource in self.vm_config.nics + self.vm_config.public_ips:
@@ -92,7 +98,8 @@ def _Create(self):
     )
 
     if self.vm_config.availability_zone:
-      cmd.extend(['--zones', self.vm_config.availability_zone])
+      zones = [vm.availability_zone for vm in self.zoned_vm_configs]
+      cmd.extend(['--zones', ','.join(zones)])
 
     if self.vm_config.boot_startup_script:
       cmd.extend(['--custom-data', self.vm_config.boot_startup_script])

perfkitbenchmarker/providers/gcp/gce_managed_instance_group.py

@@ -107,21 +107,22 @@ class GceManagedInstanceGroup(managed_vm_group.BaseManagedVmGroup):
   def __init__(
       self,
       spec: vm_group_decoders.VmGroupSpec,
-      vm_config: virtual_machine.BaseVirtualMachine,
+      vm_configs: list[virtual_machine.BaseVirtualMachine],
   ):
-    super().__init__(spec, vm_config)
+    super().__init__(spec, vm_configs)
     self.vm_config: gce_virtual_machine.GceVirtualMachine = cast(
         gce_virtual_machine.GceVirtualMachine, self.vm_config
     )
+    self.zoned_vm_configs: list[gce_virtual_machine.GceVirtualMachine] = cast(
+        list[gce_virtual_machine.GceVirtualMachine], self.zoned_vm_configs
+    )
     self.project = self.vm_config.project
     # 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
+    assert self.zones
+    if util.IsZone(self.zones[0]):
+      self.region = util.GetRegionFromZone(self.zones[0])
+    elif util.IsRegion(self.zones[0]):
+      self.region = self.zones[0]
     else:
       raise ValueError(
           f'Unsupported zone: {self.vm_config.zone}. Must be a zone or region.'
@@ -140,7 +141,7 @@ def _GcloudCmd(self, *args) -> util.GcloudCommand:
         self, 'compute', 'instance-groups', 'managed', *args
     )
     # GcloudCommand does not have support for regional resources.
-    if not self.zone:
+    if self.is_regional:
       cmd.flags['region'] = self.region
     return cmd
 
@@ -153,6 +154,9 @@ def _Create(self):
         '--size',
         str(self.vm_count),
     )
+    # --zone and --region are handled by GcloudCommand.
+    if len(self.zones) > 1:
+      cmd.args.append('--zones=' + ','.join(self.zones))
     # TODO(pclay): Consider beta and --resource-manager-tags for labels.
     cmd.Issue()
 

tests/providers/azure/azure_vm_scaling_set_test.py

@@ -69,7 +69,7 @@ def TestMig(
             os_type='debian12',
             vm_spec={'Azure': {'machine_type': 'Standard_D2s_v5'}},
         ),
-        vm_config,
+        [vm_config],
     )
 
   @mock.patch.object(

tests/providers/gcp/gce_managed_instance_group_test.py

@@ -45,7 +45,7 @@ def TestMig(self, mock_get_network, *_, zone='us-central1-c'):
             os_type='debian12',
             vm_spec={'GCP': {'machine_type': 'n1-standard-4'}},
         ),
-        vm_config,
+        [vm_config],
     )
 
   def testCreate(self, *_):
@@ -72,6 +72,40 @@ def testCreateRegional(self, *_):
         self.mock_cmd.all_commands,
     )
 
+  @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 testCreateMultiZone(self, mock_get_network, *_):
+    mock_get_network.return_value.placement_group.name = 'test_placement_group'
+    vm_configs = [
+        pkb_common_test_case.TestGceVirtualMachine(
+            gce_virtual_machine.GceVmSpec(
+                'test_component',
+                machine_type='n1-standard-4',
+                zone=zone,
+            )
+        )
+        for zone in ['us-central1-a', 'us-central1-b']
+    ]
+    mig = gce_managed_instance_group.GceManagedInstanceGroup(
+        vm_group_decoders.VmGroupSpec(
+            'test_component',
+            cloud='GCP',
+            os_type='debian12',
+            vm_spec={'GCP': {'machine_type': 'n1-standard-4'}},
+        ),
+        vm_configs,
+    )
+    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 --zones=us-central1-a,us-central1-b --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')