gce_virtual_machine: Support user provided instance metadata

Two of the instance metadata keys that are particularly interesting are
'startup-script' and 'user-data', which are reserved by google-startup-script
and Cloud
Config(https://cloudinit.readthedocs.org/en/latest/topics/format.html#cloud-config-data)
to run user specified workloads at boot time. This can be a useful way to
measure VM boot time as a lot cloud workloads are started through some init
scripts, instead of SSH. Plus, sshd on GCP images is usually started at a pretty
late stage.

perfkitbenchmarker/flag_util.py

@@ -364,3 +364,27 @@ def DEFINE_yaml(name, default, help, flag_values=flags.FLAGS, **kwargs):
   serializer = YAMLSerializer()
 
   flags.DEFINE(parser, name, default, help, flag_values, serializer, **kwargs)
+
+
+def ParseKeyValuePairs(strings):
+  """Parses colon separated key value pairs from a list of strings.
+
+  Pairs should be separated by a comma and key and value by a colon, e.g.,
+  ['k1:v1', 'k2:v2,k3:v3'].
+
+  Args:
+    strings: A list of strings.
+
+  Returns:
+    A dict populated with keys and values from the flag.
+  """
+  pairs = {}
+  for pair in [kv for s in strings for kv in s.split(',')]:
+    try:
+      key, value = pair.split(':')
+      pairs[key] = value
+    except ValueError:
+      logging.error('Bad key value pair format. Skipping "%s".', pair)
+      continue
+
+  return pairs

perfkitbenchmarker/providers/gcp/flags.py

@@ -35,3 +35,21 @@
     'https://cloud.google.com/sdk/gcloud/reference/compute/disks/create')
 flags.DEFINE_string('gce_network_name', None, 'The name of an already created '
                     'network to use instead of creating a new one.')
+flags.DEFINE_multistring(
+    'gcp_instance_metadata_from_file',
+    [],
+    'A colon separated key-value pair that will be added to the '
+    '"--metadata-from-file" flag of the gcloud cli (with the colon replaced by '
+    'the equal sign). Multiple key-value pairs may be specified by separating '
+    'each pair by commas. This option can be repeated multiple times. For '
+    'information about GCP instance metadata, see: --metadata-from-file from '
+    '`gcloud help compute instances create`.')
+flags.DEFINE_multistring(
+    'gcp_instance_metadata',
+    [],
+    'A colon separated key-value pair that will be added to the '
+    '"--metadata" flag of the gcloud cli (with the colon replaced by the equal '
+    'sign). Multiple key-value pairs may be specified by separating each pair '
+    'by commas. This option can be repeated multiple times. For information '
+    'about GCP instance metadata, see: --metadata from '
+    '`gcloud help compute instances create`.')

perfkitbenchmarker/providers/gcp/gce_virtual_machine.py

@@ -25,14 +25,17 @@
 """
 
 import json
+import logging
 import re
 import yaml
 
 from perfkitbenchmarker import disk
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import events
 from perfkitbenchmarker import flags
+from perfkitbenchmarker import flag_util
 from perfkitbenchmarker import linux_virtual_machine as linux_vm
+from perfkitbenchmarker import providers
 from perfkitbenchmarker import virtual_machine
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker import windows_virtual_machine
@@ -41,7 +44,6 @@
 from perfkitbenchmarker.providers.gcp import gce_disk
 from perfkitbenchmarker.providers.gcp import gce_network
 from perfkitbenchmarker.providers.gcp import util
-from perfkitbenchmarker import providers
 
 FLAGS = flags.FLAGS
 
@@ -283,11 +285,32 @@ def _GenerateCreateCommand(self, ssh_keys_path):
       cmd.flags['machine-type'] = self.machine_type
     cmd.flags['tags'] = 'perfkitbenchmarker'
     cmd.flags['no-restart-on-failure'] = True
-    cmd.flags['metadata-from-file'] = 'sshKeys=%s' % ssh_keys_path
-    metadata = ['owner=%s' % FLAGS.owner]
-    for key, value in self.boot_metadata.iteritems():
-      metadata.append('%s=%s' % (key, value))
-    cmd.flags['metadata'] = ','.join(metadata)
+
+    metadata_from_file = {'sshKeys': ssh_keys_path}
+    parsed_metadata_from_file = flag_util.ParseKeyValuePairs(
+        FLAGS.gcp_instance_metadata_from_file)
+    for key, value in parsed_metadata_from_file.iteritems():
+      if key in metadata_from_file:
+        logging.warning('Metadata "%s" is set internally. Cannot be overridden '
+                        'from command line.' % key)
+        continue
+      metadata_from_file[key] = value
+    cmd.flags['metadata-from-file'] = ','.join([
+        '%s=%s' % (k, v) for k, v in metadata_from_file.iteritems()
+    ])
+
+    metadata = {'owner': FLAGS.owner} if FLAGS.owner else {}
+    metadata.update(self.boot_metadata)
+    parsed_metadata = flag_util.ParseKeyValuePairs(FLAGS.gcp_instance_metadata)
+    for key, value in parsed_metadata.iteritems():
+      if key in metadata:
+        logging.warning('Metadata "%s" is set internally. Cannot be overridden '
+                        'from command line.' % key)
+        continue
+      metadata[key] = value
+    cmd.flags['metadata'] = ','.join(
+        ['%s=%s' % (k, v) for k, v in metadata.iteritems()])
+
     if not FLAGS.gce_migrate_on_maintenance:
       cmd.flags['maintenance-policy'] = 'TERMINATE'
     ssd_interface_option = NVME if NVME in self.image else SCSI

perfkitbenchmarker/publisher.py

@@ -29,6 +29,7 @@
 from perfkitbenchmarker import disk
 from perfkitbenchmarker import events
 from perfkitbenchmarker import flags
+from perfkitbenchmarker import flag_util
 from perfkitbenchmarker import version
 from perfkitbenchmarker import vm_util
 
@@ -177,14 +178,8 @@ def AddMetadata(self, metadata, benchmark_spec):
     # Flatten all user metadata into a single list (since each string in the
     # FLAGS.metadata can actually be several key-value pairs) and then iterate
     # over it.
-    for pair in [kv for item in FLAGS.metadata for kv in item.split(',')]:
-      try:
-        key, value = pair.split(':')
-        metadata[key] = value
-      except ValueError:
-        logging.error('Bad metadata flag format. Skipping "%s".', pair)
-        continue
-
+    parsed_metadata = flag_util.ParseKeyValuePairs(FLAGS.metadata)
+    metadata.update(parsed_metadata)
     return metadata
 
 

tests/gce_virtual_machine_test.py

@@ -16,6 +16,7 @@
 
 import contextlib
 import mock
+import re
 import unittest
 
 from perfkitbenchmarker import benchmark_spec
@@ -227,6 +228,8 @@ def setUp(self):
     self._mocked_flags.gcloud_path = 'test_gcloud'
     self._mocked_flags.os_type = os_types.DEBIAN
     self._mocked_flags.run_uri = 'aaaaaa'
+    self._mocked_flags.gcp_instance_metadata = []
+    self._mocked_flags.gcp_instance_metadata_from_file = []
     # Creating a VM object causes network objects to be added to the current
     # thread's benchmark spec. Create such a benchmark spec for these tests.
     self.addCleanup(context.SetThreadBenchmarkSpec, None)
@@ -267,6 +270,41 @@ def testImageProjectFlag(self):
       self.assertIn('--image-project bar',
                     ' '.join(issue_command.call_args[0][0]))
 
+  def testGcpInstanceMetadataFlag(self):
+    with self._PatchCriticalObjects() as issue_command:
+      self._mocked_flags.gcp_instance_metadata = ['k1:v1', 'k2:v2,k3:v3']
+      self._mocked_flags.owner = 'test-owner'
+      vm_spec = gce_virtual_machine.GceVmSpec(
+          'test_vm_spec.GCP', self._mocked_flags, image='image',
+          machine_type='test_machine_type')
+      vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
+      vm._Create()
+      self.assertEquals(issue_command.call_count, 1)
+      actual_metadata = re.compile('--metadata\s+(.*)(\s+--)?').search(
+          ' '.join(issue_command.call_args[0][0])).group(1)
+      self.assertIn('k1=v1', actual_metadata)
+      self.assertIn('k2=v2', actual_metadata)
+      self.assertIn('k3=v3', actual_metadata)
+      # Assert that FLAGS.owner is honored and added to instance metadata.
+      self.assertIn('owner=test-owner', actual_metadata)
+
+  def testGcpInstanceMetadataFromFileFlag(self):
+    with self._PatchCriticalObjects() as issue_command:
+      self._mocked_flags.gcp_instance_metadata_from_file = [
+          'k1:p1', 'k2:p2,k3:p3']
+      vm_spec = gce_virtual_machine.GceVmSpec(
+          'test_vm_spec.GCP', self._mocked_flags, image='image',
+          machine_type='test_machine_type')
+      vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
+      vm._Create()
+      self.assertEquals(issue_command.call_count, 1)
+      actual_metadata_from_file = re.compile(
+          '--metadata-from-file\s+(.*)(\s+--)?').search(
+          ' '.join(issue_command.call_args[0][0])).group(1)
+      self.assertIn('k1=p1', actual_metadata_from_file)
+      self.assertIn('k2=p2', actual_metadata_from_file)
+      self.assertIn('k3=p3', actual_metadata_from_file)
+
 
 if __name__ == '__main__':
   unittest.main()