def test_dimensions(self):
"""Ensures dimensions are returned."""
instance = models.Instance(key=instances.get_instance_key(
'base-name',
'revision',
'zone',
'instance-name',
), )
instance_template_revision = models.InstanceTemplateRevision(
dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
disk_size_gb=300,
machine_type='n1-standard-8',
)
expected_dimensions = {
'backend': 'GCE',
'disk_size_gb': 300,
'hostname': 'instance-name',
'memory_gb': 30,
'num_cpus': 8,
'os_family': 'LINUX',
}
self.assertEqual(
catalog.extract_dimensions(instance, instance_template_revision),
expected_dimensions,
)
def test_lease_untriaged(self):
machine_type = lease_management.MachineType(
id='untriaged',
leases=[
lease_management.MachineLease(client_request_id='fake-id-1'),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='untriaged',
target_size=2,
)
machine_type.put()
responses = rpc_to_json(
machine_provider.BatchedLeaseResponse(responses=[
machine_provider.LeaseResponse(
client_request_id='fake-id-1',
request_hash='fake-hash-1',
state=machine_provider.LeaseRequestState.UNTRIAGED,
),
]))
lease_management.update_leases(machine_type.key, responses)
updated_machine_type = machine_type.key.get()
self.assertEqual(len(updated_machine_type.leases), 1)
self.assertEqual(updated_machine_type.leases[0].client_request_id,
'fake-id-1')
self.failIf(updated_machine_type.leases[0].hostname)
self.failIf(updated_machine_type.leases[0].lease_expiration_ts)
self.assertEqual(updated_machine_type.leases[0].request_hash,
'fake-hash-1')
def test_machine_type_one_pending(self):
machine_type = lease_management.MachineType(
id='one-pending',
leases=[
lease_management.MachineLease(
client_request_id='fake-id-1',
hostname='fake-host-1',
request_hash='fake-hash-1',
),
lease_management.MachineLease(
client_request_id='fake-id-2',
hostname='fake-host-2',
request_hash='fake-hash-2',
),
lease_management.MachineLease(
client_request_id='fake-id-3',
request_hash='fake-hash-3',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='one-pending',
target_size=2,
)
requests = lease_management._generate_lease_request_status_updates(
machine_type, 'https://example.com')
self.assertEqual(len(requests), 1)
self.assertEqual(requests[0].request_id, 'fake-id-3')
def test_enable_and_update_machine_type(self):
def fetch_machine_types():
return {
'machine-type':
bots_pb2.MachineType(
early_release_secs=0,
lease_duration_secs=2,
mp_dimensions=['disk_gb:100'],
name='machine-type',
target_size=1,
),
}
self.mock(
lease_management.bot_groups_config,
'fetch_machine_types',
fetch_machine_types,
)
key = lease_management.MachineType(
id='machine-type',
early_release_secs=0,
enabled=False,
lease_duration_secs=1,
mp_dimensions=machine_provider.Dimensions(disk_gb=100, ),
target_size=1,
).put()
lease_management.ensure_entities_exist()
self.failUnless(key.get().enabled)
self.assertEqual(key.get().lease_duration_secs, 2)
def test_machine_type_no_expired_leases(self):
machine_type = lease_management.MachineType(
id='no-leases',
leases=[
lease_management.MachineLease(
client_request_id='fake-id-1',
request_hash='fake-hash-1',
),
lease_management.MachineLease(
client_request_id='fake-id-2',
hostname='fake-host',
lease_expiration_ts=(utils.utcnow() +
datetime.timedelta(seconds=60)),
request_hash='fake-hash-2',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='no-leases',
target_size=2,
)
expired = lease_management._clean_up_expired_leases(machine_type)
self.failIf(expired)
self.assertEqual(len(machine_type.leases), 2)
self.failIf(machine_type.pending_deletion)
def test_machine_type_one_expired_lease(self):
machine_type = lease_management.MachineType(
id='no-leases',
leases=[
lease_management.MachineLease(
client_request_id='fake-id-1',
request_hash='fake-hash-1',
),
lease_management.MachineLease(
client_request_id='fake-id-2',
hostname='fake-host-2',
lease_expiration_ts=utils.EPOCH,
request_hash='fake-hash-2',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
pending_deletion=[
'fake-host-3',
],
request_id_base='no-leases',
target_size=2,
)
expired = lease_management._clean_up_expired_leases(machine_type)
self.assertEqual(len(expired), 1)
self.assertEqual(expired[0], 'fake-host-2')
self.assertEqual(len(machine_type.leases), 1)
self.assertEqual(machine_type.leases[0].client_request_id, 'fake-id-1')
self.assertEqual(machine_type.leases[0].request_hash, 'fake-hash-1')
self.assertEqual(len(machine_type.pending_deletion), 2)
hostnames = sorted(machine_type.pending_deletion)
self.assertEqual(hostnames[0], 'fake-host-2')
self.assertEqual(hostnames[1], 'fake-host-3')
def test_lease_expired(self):
machine_type = lease_management.MachineType(
id='untriaged',
leases=[
lease_management.MachineLease(client_request_id='fake-id-1'),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='untriaged',
target_size=2,
)
machine_type.put()
responses = rpc_to_json(
machine_provider.BatchedLeaseResponse(responses=[
machine_provider.LeaseResponse(
client_request_id='fake-id-1',
request_hash='fake-hash-1',
state=machine_provider.LeaseRequestState.FULFILLED,
),
]))
lease_management.update_leases(machine_type.key, responses)
updated_machine_type = machine_type.key.get()
self.failIf(updated_machine_type.leases)
def test_machine_lease_exists_mismatched_updated(self):
def fetch_machine_types():
return {
'machine-type':
bots_pb2.MachineType(
early_release_secs=0,
lease_duration_secs=1,
mp_dimensions=['disk_gb:100'],
name='machine-type',
target_size=1,
),
}
self.mock(
lease_management.bot_groups_config,
'fetch_machine_types',
fetch_machine_types,
)
key = lease_management.MachineType(
id='machine-type',
early_release_secs=0,
lease_duration_secs=1,
mp_dimensions=machine_provider.Dimensions(disk_gb=100, ),
target_size=1,
).put()
key = lease_management.MachineLease(
id='%s-0' % key.id(),
early_release_secs=1,
lease_duration_secs=2,
lease_expiration_ts=utils.utcnow(),
machine_type=key,
mp_dimensions=machine_provider.Dimensions(disk_gb=200, ),
).put()
lease_management.ensure_entities_exist()
self.assertEqual(lease_management.MachineLease.query().count(), 1)
self.assertEqual(key.get().early_release_secs, 0)
self.assertEqual(key.get().lease_duration_secs, 1)
self.assertEqual(key.get().mp_dimensions.disk_gb, 100)
def test_machine_lease_exists_mismatched_not_updated(self):
key = lease_management.MachineType(
early_release_secs=0,
lease_duration_secs=1,
mp_dimensions=machine_provider.Dimensions(disk_gb=100, ),
target_size=1,
).put()
key = lease_management.MachineLease(
id='%s-0' % key.id(),
early_release_secs=1,
lease_duration_secs=2,
machine_type=key,
mp_dimensions=machine_provider.Dimensions(disk_gb=200, ),
).put()
lease_management.ensure_entities_exist()
self.assertEqual(lease_management.MachineLease.query().count(), 1)
self.assertEqual(key.get().early_release_secs, 1)
self.assertEqual(key.get().lease_duration_secs, 2)
self.assertEqual(key.get().mp_dimensions.disk_gb, 200)
def test_need_one(self):
machine_type = lease_management.MachineType(
id='need-one',
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='need-one',
target_size=1,
)
requests = lease_management._generate_lease_requests_for_new_machines(
machine_type, 'https://example.com')
self.assertEqual(len(requests), 1)
def test_machine_type_no_leases(self):
machine_type = lease_management.MachineType(
id='no-leases',
leases=[],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='no-leases',
target_size=2,
)
requests = lease_management._generate_lease_request_status_updates(
machine_type, 'https://example.com')
self.failIf(requests)
def test_machine_type_no_leases(self):
machine_type = lease_management.MachineType(
id='no-leases',
leases=[],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='no-leases',
target_size=2,
)
expired = lease_management._clean_up_expired_leases(machine_type)
self.failIf(expired)
self.failIf(machine_type.leases)
self.failIf(machine_type.pending_deletion)
def test_ensure_correct_request_ids(self):
machine_type = lease_management.MachineType(
id='ensure-correct-request-ids',
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_count=2,
request_id_base='ensure-correct-request-ids',
target_size=2,
)
requests = lease_management._generate_lease_requests_for_new_machines(
machine_type, 'https://example.com')
self.assertEqual(len(requests), 2)
request_ids = sorted(request.request_id for request in requests)
self.assertEqual(request_ids[0], 'ensure-correct-request-ids-3')
self.assertEqual(request_ids[1], 'ensure-correct-request-ids-4')
def test_machine_type_not_found(self):
machine_type = lease_management.MachineType(
id='not-found',
leases=[
lease_management.MachineLease(
client_request_id='fake-id',
request_hash='fake-hash',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='not-found',
target_size=2,
)
requests = lease_management.generate_lease_requests(
machine_type.key, 'https://example.com')
self.failIf(requests)
def test_need_one(self):
machine_type = lease_management.MachineType(
id='need-one',
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='need-one',
target_size=1,
)
machine_type.put()
requests = lease_management.generate_lease_requests(
machine_type.key, 'https://example.com')
updated_machine_type = machine_type.key.get()
self.assertEqual(len(requests), 1)
self.assertEqual(len(updated_machine_type.leases), 1)
self.assertEqual(updated_machine_type.leases[0].client_request_id,
'need-one-1')
self.assertEqual(updated_machine_type.request_count, 1)
def test_daily_schedule_resize_to_zero(self):
def fetch_machine_types():
return {
'machine-type':
bots_pb2.MachineType(
early_release_secs=0,
lease_duration_secs=1,
mp_dimensions=['disk_gb:100'],
name='machine-type',
target_size=1,
schedule=bots_pb2.Schedule(daily=[
bots_pb2.DailySchedule(
start='0:00',
end='1:00',
days_of_the_week=xrange(7),
target_size=0,
)
], ),
),
}
self.mock(
lease_management.bot_groups_config,
'fetch_machine_types',
fetch_machine_types,
)
self.mock(
lease_management.utils,
'utcnow',
lambda: datetime.datetime(1969, 1, 1, 0, 30),
)
key = lease_management.MachineType(
id='machine-type',
early_release_secs=0,
lease_duration_secs=1,
mp_dimensions=machine_provider.Dimensions(disk_gb=100, ),
target_size=1,
).put()
lease_management.ensure_entities_exist()
self.failIf(lease_management.MachineLease.query().count())
self.failIf(key.get().target_size)
def test_machine_type_none_pending(self):
machine_type = lease_management.MachineType(
id='none-pending',
leases=[
lease_management.MachineLease(
client_request_id='fake-id',
hostname='fake-host',
request_hash='fake-hash',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='none-pending',
target_size=2,
)
requests = lease_management._generate_lease_request_status_updates(
machine_type, 'https://example.com')
self.failIf(requests)
def test_leases_fulfilled(self):
machine_type = lease_management.MachineType(
id='fulfilled',
leases=[
lease_management.MachineLease(client_request_id='fake-id-1'),
lease_management.MachineLease(client_request_id='fake-id-2'),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='fulfilled',
target_size=2,
)
machine_type.put()
responses = rpc_to_json(
machine_provider.BatchedLeaseResponse(responses=[
machine_provider.LeaseResponse(
client_request_id='fake-id-1',
hostname='fake-host-1',
lease_expiration_ts=1,
request_hash='fake-hash-1',
state=machine_provider.LeaseRequestState.FULFILLED,
),
machine_provider.LeaseResponse(
client_request_id='fake-id-2',
hostname='fake-host-2',
lease_expiration_ts=2,
request_hash='fake-hash-2',
state=machine_provider.LeaseRequestState.FULFILLED,
),
]))
lease_management.update_leases(machine_type.key, responses)
updated_machine_type = machine_type.key.get()
self.assertEqual(len(updated_machine_type.leases), 2)
self.failUnless(updated_machine_type.leases[0].hostname)
self.failUnless(updated_machine_type.leases[1].hostname)
self.failUnless(updated_machine_type.leases[0].lease_expiration_ts)
self.failUnless(updated_machine_type.leases[1].lease_expiration_ts)
request_hashes = sorted(lease.request_hash
for lease in updated_machine_type.leases)
self.assertEqual(request_hashes[0], 'fake-hash-1')
self.assertEqual(request_hashes[1], 'fake-hash-2')
def test_disable_machine_type(self):
def fetch_machine_types():
return {}
self.mock(
lease_management.bot_groups_config,
'fetch_machine_types',
fetch_machine_types,
)
key = lease_management.MachineType(
id='machine-type',
early_release_secs=0,
enabled=True,
lease_duration_secs=1,
mp_dimensions=machine_provider.Dimensions(disk_gb=100, ),
target_size=1,
).put()
lease_management.ensure_entities_exist()
self.failIf(key.get().enabled)
def test_machine_type_not_enabled(self):
machine_type = lease_management.MachineType(
id='not-enabled',
enabled=False,
leases=[
lease_management.MachineLease(
client_request_id='fake-id-1',
request_hash='fake-hash',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_count=1,
request_id_base='not-found',
target_size=2,
)
machine_type.put()
requests = lease_management.generate_lease_requests(
machine_type.key, 'https://example.com')
self.assertEqual(len(requests), 1)
self.assertEqual(requests[0].request_id, 'fake-id-1')
def test_machine_type_at_capacity(self):
machine_type = lease_management.MachineType(
id='at-capacity',
leases=[
lease_management.MachineLease(
client_request_id='fake-id-1',
request_hash='fake-hash-1',
),
lease_management.MachineLease(
client_request_id='fake-id-2',
request_hash='fake-hash-2',
),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='at-capacity',
target_size=2,
)
requests = lease_management._generate_lease_requests_for_new_machines(
machine_type, 'https://example.com')
self.failIf(requests)
def test_lease_errors(self):
machine_type = lease_management.MachineType(
id='errors',
leases=[
lease_management.MachineLease(client_request_id='fake-id-1'),
lease_management.MachineLease(client_request_id='fake-id-2'),
lease_management.MachineLease(client_request_id='fake-id-3'),
],
mp_dimensions=machine_provider.Dimensions(
os_family=machine_provider.OSFamily.LINUX, ),
request_id_base='errors',
target_size=2,
)
machine_type.put()
responses = rpc_to_json(
machine_provider.BatchedLeaseResponse(responses=[
machine_provider.LeaseResponse(
client_request_id='fake-id-1',
error=machine_provider.LeaseRequestError.DEADLINE_EXCEEDED,
),
machine_provider.LeaseResponse(
client_request_id='fake-id-2',
error=machine_provider.LeaseRequestError.REQUEST_ID_REUSE,
),
machine_provider.LeaseResponse(
client_request_id='fake-id-3',
error=machine_provider.LeaseRequestError.TRANSIENT_ERROR,
),
]))
lease_management.update_leases(machine_type.key, responses)
updated_machine_type = machine_type.key.get()
self.assertEqual(len(updated_machine_type.leases), 2)
request_ids = sorted(request.client_request_id
for request in updated_machine_type.leases)
self.assertEqual(request_ids[0], 'fake-id-1')
self.assertEqual(request_ids[1], 'fake-id-3')
def get(self):
pubsub_handler = handlers_pubsub.MachineProviderSubscriptionHandler
if not pubsub_handler.is_subscribed():
logging.error(
'Pub/Sub subscription not created:\n%s',
pubsub_handler.get_subscription_name(),
)
return
# For each group manager, tell the Machine Provider about its instances.
for template in models.InstanceTemplate.query():
logging.info(
'Retrieving instance template %s from project %s',
template.template_name,
template.template_project,
)
api = gce.Project(template.template_project)
try:
instance_template = api.get_instance_template(
template.template_name)
except net.NotFoundError:
logging.error(
'Instance template does not exist: %s',
template.template_name,
)
continue
api = gce.Project(template.instance_group_project)
properties = instance_template['properties']
disk_gb = int(
properties['disks'][0]['initializeParams']['diskSizeGb'])
memory_gb = float(
gce.machine_type_to_memory(properties['machineType']))
num_cpus = gce.machine_type_to_num_cpus(properties['machineType'])
os_family = machine_provider.OSFamily.lookup_by_name(
template.os_family)
dimensions = machine_provider.Dimensions(
backend=machine_provider.Backend.GCE,
disk_gb=disk_gb,
memory_gb=memory_gb,
num_cpus=num_cpus,
os_family=os_family,
)
try:
instances = api.get_managed_instances(
template.instance_group_name, template.zone)
except net.NotFoundError:
logging.warning(
'Instance group manager does not exist: %s',
template.instance_group_name,
)
continue
policies = machine_provider.Policies(
backend_attributes=[
machine_provider.KeyValuePair(
key='group', value=template.instance_group_name),
],
backend_project=handlers_pubsub.
MachineProviderSubscriptionHandler.TOPIC_PROJECT,
backend_topic=handlers_pubsub.
MachineProviderSubscriptionHandler.TOPIC,
on_reclamation=machine_provider.MachineReclamationPolicy.
DELETE,
)
process_instance_group(
template.instance_group_name,
dimensions,
policies,
instances,
template.zone,
template.instance_group_project,
)
