def testCustomVmPreemptible(self):
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type={'cpus': 1, 'memory': '1.0GiB'},
preemptible=True)
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertEqual(vm.GetMachineTypeDict(), {'cpus': 1, 'memory_mib': 1024,
'preemptible': True})
def testCustomVmNonPreemptible(self):
spec = gce_virtual_machine.GceVmSpec(_COMPONENT, machine_type={
'cpus': 1, 'memory': '1.0GiB'}, project='p')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertDictContainsSubset(
{'cpus': 1, 'memory_mib': 1024, 'project': 'p',
'dedicated_host': False},
vm.GetResourceMetadata())
def testMigrateOnMaintenanceFlagFalseWithGpus(self):
with PatchCriticalObjects() as issue_command:
self._mocked_flags['gce_migrate_on_maintenance'].parse(False)
vm_spec = gce_virtual_machine.GceVmSpec(
'test_vm_spec.GCP', self._mocked_flags, image='image',
machine_type='test_machine_type', gpu_count=1, gpu_type='k80')
vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
vm._Create()
self.assertEquals(issue_command.call_count, 1)
def testVmWithMachineTypeNonPreemptible(self):
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type='test_machine_type', project='p')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertDictContainsSubset(
{'dedicated_host': False, 'machine_type': 'test_machine_type',
'project': 'p'},
vm.GetResourceMetadata()
)
def testVmWithMachineTypePreemptible(self):
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type='test_machine_type',
preemptible=True)
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertEqual(vm.GetMachineTypeDict(), {
'machine_type': 'test_machine_type',
'preemptible': True
})
def testPreemptibleVMFlag(self):
with self._PatchCriticalObjects() as issue_command:
self._mocked_flags['gce_preemptible_vms'].Parse(True)
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)
self.assertIn('--preemptible', issue_command.call_args[0][0])
def testVmWithoutGpu(self):
with PatchCriticalObjects() as issue_command:
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type={
'cpus': 1,
'memory': '1.0GiB',
})
vm = gce_virtual_machine.GceVirtualMachine(spec)
vm._Create()
self.assertEqual(issue_command.call_count, 1)
self.assertNotIn('--accelerator', issue_command.call_args[0][0])
def testCreateVMAlreadyExists(self):
fake_rets = [('stdout', 'The resource already exists', 1)]
with PatchCriticalObjects(fake_rets):
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type={
'cpus': 1,
'memory': '1.0GiB',
})
vm = gce_virtual_machine.GceVirtualMachine(spec)
with self.assertRaises(errors.Resource.CreationError):
vm._Create()
def testVmWithMachineTypeNonPreemptible(self):
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type='test_machine_type',
project='p')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertEqual(
vm.GetMachineTypeDict(), {
'dedicated_host': False,
'machine_type': 'test_machine_type',
'project': 'p'
})
def testAcceleratorTypeOverrideFlag(self):
with PatchCriticalObjects() as issue_command:
self._mocked_flags['gce_accelerator_type_override'].parse('fake_type')
vm_spec = gce_virtual_machine.GceVmSpec(
'test_vm_spec.GCP', self._mocked_flags, image='image',
machine_type='test_machine_type', gpu_count=1, gpu_type='k80')
vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
vm._Create()
self.assertEquals(issue_command.call_count, 1)
self.assertIn('--accelerator', issue_command.call_args[0][0])
self.assertIn('type=fake_type,count=1',
issue_command.call_args[0][0])
def testCustomVmWithGpus(self):
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT,
machine_type={'cpus': 1, 'memory': '1.0GiB'},
gpu_count=2,
gpu_type='k80',
project='fakeproject')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertDictContainsSubset({
'cpus': 1, 'memory_mib': 1024, 'project': 'fakeproject',
'dedicated_host': False, 'gpu_count': 2, 'gpu_type': 'k80'
}, vm.GetResourceMetadata())
def testImageProjectFlag(self):
"""Tests that custom image_project flag is supported."""
with self._PatchCriticalObjects() as issue_command:
self._mocked_flags.image_project = 'bar'
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)
self.assertIn('--image-project bar',
' '.join(issue_command.call_args[0][0]))
def _CreateVmCommand(self, **flag_kwargs):
with PatchCriticalObjects() as issue_command:
for key, value in flag_kwargs.items():
FLAGS[key].parse(value)
vm_spec = gce_virtual_machine.GceVmSpec(
'test_vm_spec.GCP',
FLAGS,
image='image',
machine_type='test_machine_type')
vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
vm._Create()
return ' '.join(issue_command.call_args[0][0]), issue_command.call_count
def testMigrateOnMaintenanceFlagTrueWithGpus(self):
with PatchCriticalObjects():
self._mocked_flags['gce_migrate_on_maintenance'].parse(True)
vm_spec = gce_virtual_machine.GceVmSpec(
'test_vm_spec.GCP', self._mocked_flags, image='image',
machine_type='test_machine_type', gpu_count=1, gpu_type='k80')
vm = gce_virtual_machine.GceVirtualMachine(vm_spec)
with self.assertRaises(errors.Config.InvalidValue) as cm:
vm._Create()
self.assertEqual(str(cm.exception), (
'Cannot set flag gce_migrate_on_maintenance on '
'instances with GPUs, as it is not supported by GCP.'))
def testCustomVmNonPreemptible(self):
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type={
'cpus': 1,
'memory': '1.0GiB'
},
project='p')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertEqual(vm.GetMachineTypeDict(), {
'cpus': 1,
'memory_mib': 1024,
'project': 'p',
'dedicated_host': False
})
def testVmWithGpu(self):
with PatchCriticalObjects() as issue_command:
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT,
machine_type='n1-standard-8',
gpu_count=2,
gpu_type='k80')
vm = gce_virtual_machine.GceVirtualMachine(spec)
vm._Create()
self.assertEqual(issue_command.call_count, 1)
self.assertIn('--accelerator', issue_command.call_args[0][0])
self.assertIn('type=nvidia-tesla-k80,count=2',
issue_command.call_args[0][0])
self.assertIn('--maintenance-policy', issue_command.call_args[0][0])
self.assertIn('TERMINATE', issue_command.call_args[0][0])
def testCreateRateLimitedMachineCreated(self, mock_cmd):
fake_rets = [('stdout', 'Rate Limit Exceeded', 1),
('stdout', 'Rate Limit Exceeded', 1),
('stdout', 'Rate Limit Exceeded', 1),
('stdout', 'Rate Limit Exceeded', 1),
('stdout', 'The resource already exists', 1)]
with PatchCriticalObjects(fake_rets) as issue_command:
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type={
'cpus': 1,
'memory': '1.0GiB',
})
vm = gce_virtual_machine.GceVirtualMachine(spec)
vm._Create() # No error should be thrown.
self.assertEqual(issue_command.call_count, 5)
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)
def testCreateRateLimitedMachineCreatedFailure(self, mock_cmd):
fake_rets = []
for _ in range(0, 100):
fake_rets.append(('stdout', 'Rate Limit Exceeded', 1))
with PatchCriticalObjects(fake_rets) as issue_command:
spec = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type={
'cpus': 1,
'memory': '1.0GiB',
})
vm = gce_virtual_machine.GceVirtualMachine(spec)
with self.assertRaises(
errors.Benchmarks.QuotaFailure.RateLimitExceededError):
vm._Create()
self.assertEqual(issue_command.call_count,
util.RATE_LIMITED_MAX_RETRIES + 1)
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 testCustomVmWithGpus(self):
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type={
'cpus': 1,
'memory': '1.0GiB'
},
gpu_count=2,
gpu_type='k80',
project='fakeproject')
vm = gce_virtual_machine.GceVirtualMachine(spec)
self.assertDictEqual(
vm.GetMachineTypeDict(), {
'cpus': 1,
'memory_mib': 1024,
'project': 'fakeproject',
'dedicated_host': False,
'gpu_count': 2,
'gpu_type': 'k80'
})
