def update_mappings_for_fuzzer(fuzzer, mappings=None):
"""Clear existing mappings for a fuzzer, and replace them."""
if mappings is None:
mappings = fuzzer.jobs
query = data_types.FuzzerJob.query()
query = query.filter(data_types.FuzzerJob.fuzzer == fuzzer.name)
entities = ndb_utils.get_all_from_query(query)
old_mappings = {}
for entity in entities:
old_mappings[entity.job] = entity
new_mappings = []
for job_name in mappings:
mapping = old_mappings.pop(job_name, None)
if mapping:
continue
job = data_types.Job.query(data_types.Job.name == job_name).get()
if not job:
logs.log_error('An unknown job %s was selected for fuzzer %s.' %
(job_name, fuzzer.name))
continue
mapping = data_types.FuzzerJob()
mapping.fuzzer = fuzzer.name
mapping.job = job_name
mapping.platform = job.platform
new_mappings.append(mapping)
ndb_utils.put_multi(new_mappings)
ndb_utils.delete_multi([m.key for m in list(old_mappings.values())])
def _query_and_upload_strategy_probabilities(engine):
"""Uploads queried data into datastore.
Calls query functions and uploads query results
to datastore to use as new probabilities. Probabilities
are based on new_edges feature."""
strategy_data = []
data = _query_multi_armed_bandit_probabilities(engine)
logs.log('Queried distribution for {}.'.format(engine.name))
# TODO(mukundv): Update once we choose a temperature parameter for final
# implementation.
for row in data:
curr_strategy = data_types.FuzzStrategyProbability()
curr_strategy.strategy_name = str(row['strategy'])
curr_strategy.probability = float(row['bandit_weight'])
curr_strategy.engine = engine.name
strategy_data.append(curr_strategy)
query = data_types.FuzzStrategyProbability.query(
data_types.FuzzStrategyProbability.engine == engine.name)
ndb_utils.delete_multi(
[entity.key for entity in ndb_utils.get_all_from_query(query)])
ndb_utils.put_multi(strategy_data)
logs.log('Uploaded queried distribution to ndb for {}'.format(engine.name))
_store_probabilities_in_bigquery(engine, data)
logs.log('Uploaded queried distribution to BigQuery for {}'.format(
engine.name))
def _process_project(project, bucket):
"""Collects coverage information for all fuzz targets in the given project and
the total stats for the project."""
project_name = _basename(project)
logs.log('Processing coverage for %s project.' % project_name)
report_path = storage.get_cloud_storage_file_path(bucket, project)
report_info = _read_json(report_path)
if not report_info:
logs.log_warn('Skipping code coverage for %s project.' % project_name)
return
# Iterate through report_info['fuzzer_stats_dir'] and prepare
# CoverageInformation entities for invididual fuzz targets.
entities = []
for fuzzer in storage.list_blobs(report_info['fuzzer_stats_dir'],
recursive=False):
entities.append(
_process_fuzzer_stats(fuzzer, report_info, project_name, bucket))
logs.log('Processed coverage for %d targets in %s project.' %
(len(entities), project_name))
# Prepare CoverageInformation entity for the total project stats.
entities.append(_process_project_stats(report_info, project_name))
ndb_utils.put_multi(entities)
def execute(args):
"""Build keywords."""
count_diff = 0
query = data_types.Testcase.query().order(-data_types.Testcase.timestamp)
for testcases in batcher.iterate(query, BATCH_SIZE):
for testcase in testcases:
before_testcase = to_dict(testcase)
attribute_builder.populate(testcase)
after_testcase = to_dict(testcase)
diff = get_diff(before_testcase, after_testcase)
if (count_diff % 10) == 0 and diff:
print('Migrate (dry=%s) id:%s\n%s' %
(not args.non_dry_run, testcase.key.id(), diff))
if diff:
count_diff += 1
if args.non_dry_run:
try:
ndb_utils.put_multi(testcases)
except Exception:
for testcase in testcases:
try:
testcase.put()
except Exception:
print('Error: %s %s' %
(testcase.key.id(), sys.exc_info()))
print('Done (count_diff=%d)' % count_diff)
def update_platform_for_job(job_name, new_platform):
"""Update platform for all mappings for a particular job."""
query = data_types.FuzzerJob.query()
query = query.filter(data_types.FuzzerJob.job == job_name)
mappings = ndb_utils.get_all_from_query(query)
new_mappings = []
for mapping in mappings:
mapping.platform = new_platform
new_mappings.append(mapping)
ndb_utils.put_multi(new_mappings)
def assign_hosts_to_workers(self):
"""Assign host instances to workers."""
all_host_names = set()
for assignment in self.gce_project.host_worker_assignments:
host_cluster = self.gce_project.get_cluster(assignment.host)
worker_cluster = self.gce_project.get_cluster(assignment.worker)
if host_cluster.gce_zone != worker_cluster.gce_zone:
logging.error('Mismatching zones for %s and %s.',
assignment.host, assignment.worker)
continue
if (host_cluster.instance_count * assignment.workers_per_host !=
worker_cluster.instance_count):
logging.error(
'Invalid host/worker cluster size for %s and %s.',
assignment.host, assignment.worker)
continue
if host_cluster.high_end != worker_cluster.high_end:
logging.error('Mismatching high end setting for %s and %s',
assignment.host, assignment.worker)
continue
manager = bot_manager.BotManager(self.gce_project.project_id,
host_cluster.gce_zone)
host_instance_group = manager.instance_group(host_cluster.name)
if not host_instance_group.exists():
logging.error('Host instance group %s does not exist.',
host_cluster.name)
continue
host_names = [
_instance_name_from_url(instance['instance'])
for instance in host_instance_group.list_managed_instances()
]
all_host_names.update(host_names)
worker_instances = self.get_all_workers_in_cluster(
manager, worker_cluster.name)
if len(worker_instances) != worker_cluster.instance_count:
logging.error(
'Actual number of worker instances for %s did not match. '
'Expected %d, got %d.', worker_cluster.name,
worker_cluster.instance_count, len(worker_instances))
continue
new_assignments = self.do_assign_hosts_to_workers(
host_names, worker_instances, assignment.workers_per_host)
ndb_utils.put_multi(new_assignments)
self.cleanup_old_assignments(all_host_names)
def update_job_weight(job_name, multiplier):
"""Update a job weight."""
tool_name = environment.get_memory_tool_name(job_name)
multiplier *= SANITIZER_WEIGHTS.get(tool_name, DEFAULT_SANITIZER_WEIGHT)
engine = environment.get_engine_for_job(job_name)
multiplier *= ENGINE_WEIGHTS.get(engine, DEFAULT_ENGINE_WEIGHT)
query = data_types.FuzzerJob.query(data_types.FuzzerJob.job == job_name)
changed_weights = []
for fuzzer_job in query:
if fuzzer_job.multiplier != multiplier:
fuzzer_job.multiplier = multiplier
changed_weights.append(fuzzer_job)
if changed_weights:
ndb_utils.put_multi(changed_weights)
def update_admins(new_admins):
"""Update list of admins."""
existing_admins = ndb_utils.get_all_from_model(data_types.Admin)
to_remove = []
existing_admin_emails = set()
for admin in existing_admins:
if admin.email not in new_admins:
logs.log('Removing admin ' + admin.email)
to_remove.append(admin.key)
existing_admin_emails.add(admin.email)
ndb_utils.delete_multi(to_remove)
to_add = []
for admin in new_admins:
if admin not in existing_admin_emails:
to_add.append(data_types.Admin(id=admin, email=admin))
logs.log('Adding admin ' + admin)
ndb_utils.put_multi(to_add)
def update_project_cpus(self):
"""Update CPU allocations for each project."""
all_projects = list(data_types.OssFuzzProject.query().order(
data_types.OssFuzzProject.name))
self.cleanup_old_projects([project.name for project in all_projects])
projects = [
project for project in all_projects if not project.high_end
]
high_end_projects = [
project for project in all_projects if project.high_end
]
project_infos = [
self.get_or_create_project_info(project.name)
for project in projects
]
high_end_project_infos = [
self.get_or_create_project_info(project.name)
for project in high_end_projects
]
for project, project_info in itertools.chain(
list(zip(projects, project_infos)),
list(zip(high_end_projects, high_end_project_infos))):
self.cleanup_clusters(project, project_info)
self.start_thread_pool()
# Calculate CPUs in each cluster.
for cluster in self.gce_project.clusters:
if not cluster.distribute:
self.pending_updates.append(
self.thread_pool.submit(self.update_cluster, cluster,
cluster.name,
cluster.instance_count))
continue
if cluster.high_end:
current_projects = high_end_projects
current_project_infos = high_end_project_infos
else:
current_projects = projects
current_project_infos = project_infos
cpu_counts = self.distribute_cpus(current_projects,
cluster.instance_count)
# Resize projects starting with ones that reduce number of CPUs. This is
# so that we always have quota when we're resizing a project cluster.
# pylint: disable=cell-var-from-loop
def _cpu_diff_key(index):
cluster_info = current_project_infos[index].get_cluster_info(
cluster.name)
if cluster_info and cluster_info.cpu_count is not None:
old_cpu_count = cluster_info.cpu_count
else:
old_cpu_count = 0
return cpu_counts[index] - old_cpu_count
resize_order = sorted(list(range(len(cpu_counts))),
key=_cpu_diff_key)
for i in resize_order:
project = current_projects[i]
project_info = current_project_infos[i]
self.update_project_cluster(project,
project_info,
cluster,
cpu_counts[i],
disk_size_gb=project.disk_size_gb)
self.finish_updates()
ndb_utils.put_multi(project_infos)
ndb_utils.put_multi(high_end_project_infos)
