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, )