def __init__(self, cluster_spec):
super().__init__(user_managed=bool(cluster_spec.static_cluster))
self.name = cluster_spec.static_cluster or 'pkb-' + FLAGS.run_uri
self.vm_config = virtual_machine.GetVmClass(self.CLOUD, os_types.DEFAULT)(
cluster_spec.vm_spec)
self.zone = self.vm_config.zone
# Use Virtual Machine class to resolve VM Spec. This lets subclasses parse
# Provider specific information like disks out of the spec.
for name, nodepool in cluster_spec.nodepools.copy().items():
nodepool_zone = nodepool.vm_spec.zone
# VM Classes can require zones. But nodepools have optional zones.
if not nodepool_zone:
nodepool.vm_spec.zone = self.zone
nodepool.vm_config = virtual_machine.GetVmClass(
self.CLOUD, os_types.DEFAULT)(nodepool.vm_spec)
nodepool.vm_config.zone = nodepool_zone
nodepool.num_nodes = nodepool.vm_count
# Fix name
del cluster_spec.nodepools[name]
cluster_spec.nodepools[NodePoolName(name)] = nodepool
self.nodepools = cluster_spec.nodepools
self.num_nodes = cluster_spec.vm_count
self.min_nodes = cluster_spec.min_vm_count or self.num_nodes
self.max_nodes = cluster_spec.max_vm_count or self.num_nodes
self.containers = collections.defaultdict(list)
self.services = {}
def testCosVm(self):
vm_class = virtual_machine.GetVmClass(providers.GCP, os_types.COS)
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type='fake-machine-type')
fake_image = 'fake_cos_image'
with PatchCriticalObjects(
self._CreateFakeReturnValues(fake_image)) as issue_command:
vm = vm_class(spec)
vm._Create()
vm.created = True
command_string = ' '.join(issue_command.call_args[0][0])
self.assertEqual(issue_command.call_count, 1)
self.assertIn('gcloud compute instances create', command_string)
self.assertIn('--image-family cos-stable', command_string)
self.assertIn('--image-project cos-cloud', command_string)
vm._PostCreate()
self.assertEqual(issue_command.call_count, 3)
vm_metadata = vm.GetResourceMetadata()
self.assertDictContainsSubset(
{
'image': fake_image,
'image_family': 'cos-stable',
'image_project': 'cos-cloud'
}, vm_metadata)
def testCreateUbuntuInCustomDisk(self):
"""Test simulating passing --image and --image_project."""
vm_class = virtual_machine.GetVmClass(providers.GCP,
os_types.UBUNTU1604)
fake_image = 'fake-ubuntu1604'
fake_image_project = 'fake-project'
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type='fake-machine-type',
image=fake_image,
image_project=fake_image_project,
boot_disk_size=20,
boot_disk_type='fake-disk-type')
with PatchCriticalObjects(
self._CreateFakeReturnValues(fake_image)) as issue_command:
vm = vm_class(spec)
vm._Create()
vm.created = True
command_string = ' '.join(issue_command.call_args[0][0])
self.assertEqual(issue_command.call_count, 1)
self.assertIn('gcloud compute instances create', command_string)
self.assertIn('--boot-disk-size 20', command_string)
self.assertIn('--boot-disk-type fake-disk-type', command_string)
vm._PostCreate()
self.assertEqual(issue_command.call_count, 2)
vm_metadata = vm.GetResourceMetadata()
self.assertDictContainsSubset(
{
'image': fake_image,
'image_project': 'fake-project',
'boot_disk_size': 20,
'boot_disk_type': 'fake-disk-type'
}, vm_metadata)
self.assertNotIn('image_family', vm_metadata)
def testCreateUbuntu1604(self):
vm_class = virtual_machine.GetVmClass(providers.GCP,
os_types.UBUNTU1604)
fake_image = 'fake-ubuntu1604'
with PatchCriticalObjects(
self._CreateFakeReturnValues(fake_image)) as issue_command:
vm = vm_class(self.spec)
vm._Create()
vm.created = True
command_string = ' '.join(issue_command.call_args[0][0])
self.assertEqual(issue_command.call_count, 1)
self.assertIn('gcloud compute instances create', command_string)
self.assertIn(
'--image-family ubuntu-1604-lts --image-project ubuntu-os-cloud',
command_string)
self.assertNotIn('--boot-disk-size', command_string)
self.assertNotIn('--boot-disk-type', command_string)
vm._PostCreate()
self.assertEqual(issue_command.call_count, 3)
self.assertDictContainsSubset(
{
'image': fake_image,
'image_family': 'ubuntu-1604-lts',
'image_project': 'ubuntu-os-cloud',
'boot_disk_size': '10',
'boot_disk_type': 'pd-standard'
}, vm.GetResourceMetadata())
def testCreateCentos7CustomImage(self):
vm_class = virtual_machine.GetVmClass(providers.GCP, os_types.CENTOS7)
fake_image = 'fake-custom-centos7-image'
fake_image_project = 'fake-project'
spec = gce_virtual_machine.GceVmSpec(_COMPONENT,
machine_type='fake-machine-type',
image=fake_image,
image_project=fake_image_project)
with PatchCriticalObjects(self._CreateFakeReturnValues()) as issue_command:
vm = vm_class(spec)
vm._Create()
vm.created = True
command_string = ' '.join(issue_command.call_args[0][0])
self.assertEqual(issue_command.call_count, 1)
self.assertIn('gcloud compute instances create', command_string)
self.assertIn('--image ' + fake_image, command_string)
self.assertIn('--image-project ' + fake_image_project, command_string)
vm._PostCreate()
self.assertEqual(issue_command.call_count, 2)
vm_metadata = vm.GetResourceMetadata()
self.assertDictContainsSubset({'image': fake_image,
'image_project': fake_image_project},
vm_metadata)
self.assertNotIn('image_family', vm_metadata)
def create_kubernetes_vm(os_type):
spec = kubernetes_pod_spec.KubernetesPodSpec(
_COMPONENT)
vm_class = virtual_machine.GetVmClass(providers.KUBERNETES,
os_type)
kub_vm = vm_class(spec)
kub_vm._WaitForPodBootCompletion = lambda: None
kub_vm._Create()
def __init__(self, cluster_spec):
super(BaseContainerCluster, self).__init__()
self.name = 'pkb-%s' % FLAGS.run_uri
# Use Virtual Machine class to resolve VM Spec. This lets subclasses parse
# Provider specific information like disks out of the spec.
for name, nodepool in cluster_spec.nodepools.copy().items():
nodepool.vm_config = virtual_machine.GetVmClass(
self.CLOUD, os_types.DEFAULT)(nodepool.vm_spec)
nodepool.num_nodes = nodepool.vm_count
# Fix name
del cluster_spec.nodepools[name]
cluster_spec.nodepools[NodePoolName(name)] = nodepool
self.nodepools = cluster_spec.nodepools
self.vm_config = virtual_machine.GetVmClass(
self.CLOUD, os_types.DEFAULT)(cluster_spec.vm_spec)
self.num_nodes = cluster_spec.vm_count
self.min_nodes = cluster_spec.min_vm_count or self.num_nodes
self.max_nodes = cluster_spec.max_vm_count or self.num_nodes
self.containers = collections.defaultdict(list)
self.services = {}
self.zone = self.vm_config.zone
def testCreateDebian(self):
vm_class = virtual_machine.GetVmClass(providers.GCP, os_types.DEBIAN)
with PatchCriticalObjects(
self._CreateFakeReturnValues()) as issue_command:
vm = vm_class(self.spec)
vm._Create()
command_string = ' '.join(issue_command.call_args[0][0])
self.assertEqual(issue_command.call_count, 1)
self.assertIn('gcloud compute instances create', command_string)
self.assertIn('--image ubuntu-14-04', command_string)
vm._PostCreate()
self.assertEqual(issue_command.call_count, 2)
self.assertDictContainsSubset({'image': 'ubuntu-14-04'},
vm.GetResourceMetadata())
def testCreatePodBodyWrittenCorrectly(self):
spec = self.create_virtual_machine_spec()
vm_class = virtual_machine.GetVmClass(providers.KUBERNETES,
os_types.UBUNTU1604_CUDA9)
with patch_critical_objects() as (_, temp_file):
kub_vm = vm_class(spec)
# Need to set the name explicitly on the instance because the test
# running is currently using a single PKB instance, so the BaseVm
# instance counter is at an unpredictable number at this stage, and it is
# used to set the name.
kub_vm.name = _NAME
kub_vm._WaitForPodBootCompletion = lambda: None
kub_vm._Create()
write_mock = get_write_mock_from_temp_file_mock(temp_file)
self.assertJsonEqual(write_mock.call_args[0][0],
_EXPECTED_CALL_BODY_WITH_NVIDIA_CUDA_IMAGE)
def _CreateVirtualMachine(self, vm_spec, os_type, cloud):
"""Create a vm in zone.
Args:
vm_spec: A virtual_machine.BaseVmSpec object.
os_type: The type of operating system for the VM. See the flag of the
same name for more information.
cloud: The cloud for the VM. See the flag of the same name for more
information.
Returns:
A virtual_machine.BaseVirtualMachine object.
"""
vm = static_vm.StaticVirtualMachine.GetStaticVirtualMachine()
if vm:
return vm
vm_class = virtual_machine.GetVmClass(cloud, os_type)
if vm_class is None:
raise errors.Error(
'VMs of type %s" are not currently supported on cloud "%s".' %
(os_type, cloud))
return vm_class(vm_spec)