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)
