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' })