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 = pkb_common_test_case.TestGceVirtualMachine(vm_spec) vm._Create() return ' '.join( issue_command.call_args[0][0]), issue_command.call_count
def testGenerateCreateCommand(self, machine_type, min_cpu_platform_flag, min_cpu_platform_in_command): spec = gce_virtual_machine.GceVmSpec( _COMPONENT, machine_type=machine_type, gpu_count=2, gpu_type='t4', project='fakeproject', min_cpu_platform=min_cpu_platform_flag) vm = pkb_common_test_case.TestGceVirtualMachine(spec) gcloud_cmd = vm._GenerateCreateCommand('x') self.assertEqual(min_cpu_platform_in_command, gcloud_cmd.flags.get('min-cpu-platform'))
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 = pkb_common_test_case.TestGceVirtualMachine(spec) vm.created = True self.assertDictContainsSubset({ 'cpus': 1, 'memory_mib': 1024, 'project': 'fakeproject', 'dedicated_host': False, 'gpu_count': 2, 'gpu_type': 'k80' }, vm.GetResourceMetadata())
def testVmWithMachineTypePreemptible(self): spec = gce_virtual_machine.GceVmSpec(_COMPONENT, machine_type='test_machine_type', preemptible=True, project='p') vm = pkb_common_test_case.TestGceVirtualMachine(spec) vm.created = True self.assertDictContainsSubset( { 'dedicated_host': False, 'machine_type': 'test_machine_type', 'preemptible': True, 'project': 'p' }, vm.GetResourceMetadata())
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 = pkb_common_test_case.TestGceVirtualMachine(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 = pkb_common_test_case.TestGceVirtualMachine(spec) vm._Create() # No error should be thrown. self.assertEqual(issue_command.call_count, 5)
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 = pkb_common_test_case.TestGceVirtualMachine(spec) with self.assertRaises( errors.Benchmarks.QuotaFailure.RateLimitExceededError): vm._Create() self.assertEqual(issue_command.call_count, util.RATE_LIMITED_MAX_RETRIES + 1)
def testCustomVmNonPreemptible(self): spec = gce_virtual_machine.GceVmSpec(_COMPONENT, machine_type={ 'cpus': 1, 'memory': '1.0GiB' }, project='p') vm = pkb_common_test_case.TestGceVirtualMachine(spec) vm.created = True self.assertDictContainsSubset( { 'cpus': 1, 'memory_mib': 1024, 'project': 'p', 'dedicated_host': False }, vm.GetResourceMetadata())
def testUpdateInterruptibleVmStatus(self): spec = gce_virtual_machine.GceVmSpec(_COMPONENT, machine_type='test_machine_type', preemptible=True, project='p') vm = pkb_common_test_case.TestGceVirtualMachine(spec) vm.created = True vm.name = 'pkb-1234-0' vm.zone = 'asia-northeast2-a' stdout = """ [{ "endTime": "2021-03-12T10:38:00.048-08:00", "id": "5583038217590279140", "insertTime": "2021-03-12T10:38:00.048-08:00", "kind": "compute#operation", "name": "systemevent-1615574280048-5bd5b33121158-c654ffc9-0ff0dc28", "operationType": "compute.instances.preempted", "progress": 100, "selfLink": "https://www.googleapis.com/compute/v1/projects/my-project/zones/asia-northeast2-a/operations/systemevent-1615574280048-5bd5b33121158-c654ffc9-0ff0dc28", "startTime": "2021-03-12T10:38:00.048-08:00", "status": "DONE", "statusMessage": "Instance was preempted.", "targetId": "3533220210345905302", "targetLink": "https://www.googleapis.com/compute/v1/projects/my-project/zones/asia-northeast2-a/instances/pkb-1234-0", "user": "******", "zone": "https://www.googleapis.com/compute/v1/projects/my-project/zones/asia-northeast2-a" }] """ fake_rets = [(stdout, 'stderr', 0)] with PatchCriticalObjects(fake_rets) as issue_command: vm._UpdateInterruptibleVmStatusThroughApi() command_string = ' '.join(issue_command.call_args[0][0]) self.assertRegex( command_string, 'gcloud compute operations list --filter' fr' targetLink.scope\(\):{vm.name} --format json ' f'--project p --quiet --zones {vm.zone}') self.assertTrue(vm.spot_early_termination)