123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670 |
- # Copyright 2020 The Chromium Authors. All rights reserved.
- # Use of this source code is governed by a BSD-style license that can be
- # found in the LICENSE file.
- """Methods related to querying the ResultDB BigQuery tables."""
- import json
- import logging
- import math
- import multiprocessing.pool
- import os
- import subprocess
- import threading
- import time
- from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
- import six
- from typ import expectations_parser
- from typ import json_results
- from unexpected_passes_common import builders as builders_module
- from unexpected_passes_common import constants
- from unexpected_passes_common import data_types
- from unexpected_passes_common import multiprocessing_utils
- DEFAULT_NUM_SAMPLES = 100
- MAX_ROWS = (2**31) - 1
- MAX_QUERY_TRIES = 3
- # Used to prevent us from triggering too many queries simultaneously and causing
- # a bunch of rate limit errors. Anything below 1.5 seemed to result in enough
- # rate limit errors to cause problems. Raising above that for safety.
- QUERY_DELAY = 2
- # The target number of results/rows per query when running in large query mode.
- # Higher values = longer individual query times and higher chances of running
- # out of memory in BigQuery. Lower values = more parallelization overhead and
- # more issues with rate limit errors.
- TARGET_RESULTS_PER_QUERY = 20000
- # Subquery for getting all try builds that were used for CL submission. 30 days
- # is chosen because the ResultDB tables we pull data from only keep data around
- # for 30 days.
- SUBMITTED_BUILDS_TEMPLATE = """\
- SELECT
- CONCAT("build-", CAST(unnested_builds.id AS STRING)) as id
- FROM
- `commit-queue.{project_view}.attempts`,
- UNNEST(builds) as unnested_builds,
- UNNEST(gerrit_changes) as unnested_changes
- WHERE
- unnested_builds.host = "cr-buildbucket.appspot.com"
- AND unnested_changes.submit_status = "SUCCESS"
- AND start_time > TIMESTAMP_SUB(CURRENT_TIMESTAMP(),
- INTERVAL 30 DAY)"""
- QueryResult = Dict[str, Any]
- QueryParameters = Dict[str, Dict[str, Any]]
- # pylint: disable=super-with-arguments,useless-object-inheritance
- class BigQueryQuerier(object):
- """Class to handle all BigQuery queries for a script invocation."""
- def __init__(self, suite: Optional[str], project: str, num_samples: int,
- large_query_mode: bool):
- """
- Args:
- suite: A string containing the name of the suite that is being queried
- for. Can be None if there is no differentiation between different
- suites.
- project: A string containing the billing project to use for BigQuery.
- num_samples: An integer containing the number of builds to pull results
- from.
- large_query_mode: A boolean indicating whether large query mode should be
- used. In this mode, an initial, smaller query is made and its results
- are used to perform additional filtering on a second, larger query in
- BigQuery. This works around hitting a hard memory limit when running
- the ORDER BY clause.
- """
- self._suite = suite
- self._project = project
- self._num_samples = num_samples or DEFAULT_NUM_SAMPLES
- self._large_query_mode = large_query_mode
- assert self._num_samples > 0
- def FillExpectationMapForBuilders(
- self, expectation_map: data_types.TestExpectationMap,
- builders: Iterable[data_types.BuilderEntry]
- ) -> Dict[str, data_types.ResultListType]:
- """Fills |expectation_map| with results from |builders|.
- Args:
- expectation_map: A data_types.TestExpectationMap. Will be modified
- in-place.
- builders: An iterable of data_types.BuilderEntry containing the builders
- to query.
- Returns:
- A dict containing any results that were retrieved that did not have a
- matching expectation in |expectation_map| in the following format:
- {
- |builder_type|:|builder_name| (str): [
- result1 (data_types.Result),
- result2 (data_types.Result),
- ...
- ],
- }
- """
- assert isinstance(expectation_map, data_types.TestExpectationMap)
- # Ensure that all the builders are of the same type since we make some
- # assumptions about that later on.
- assert builders
- builder_type = None
- for b in builders:
- if builder_type is None:
- builder_type = b.builder_type
- else:
- assert b.builder_type == builder_type
- # Filter out any builders that we can easily determine do not currently
- # produce data we care about.
- builders = self._FilterOutInactiveBuilders(builders, builder_type)
- # Spin up a separate process for each query/add step. This is wasteful in
- # the sense that we'll have a bunch of idle processes once faster steps
- # start finishing, but ensures that we start slow queries early and avoids
- # the overhead of passing large amounts of data between processes. See
- # crbug.com/1182459 for more information on performance considerations.
- process_pool = multiprocessing_utils.GetProcessPool(nodes=len(builders))
- args = [(b, expectation_map) for b in builders]
- results = process_pool.map(self._QueryAddCombined, args)
- tmp_expectation_map = data_types.TestExpectationMap()
- all_unmatched_results = {}
- for (unmatched_results, prefixed_builder_name, merge_map) in results:
- tmp_expectation_map.Merge(merge_map, expectation_map)
- if unmatched_results:
- all_unmatched_results[prefixed_builder_name] = unmatched_results
- expectation_map.clear()
- expectation_map.update(tmp_expectation_map)
- return all_unmatched_results
- def _FilterOutInactiveBuilders(self,
- builders: Iterable[data_types.BuilderEntry],
- builder_type: str
- ) -> List[data_types.BuilderEntry]:
- """Filters out any builders that are not producing data.
- This helps save time on querying, as querying for the builder names is cheap
- while querying for individual results from a builder is expensive. Filtering
- out inactive builders lets us preemptively remove builders that we know we
- won't get any data from, and thus don't need to waste time querying.
- Args:
- builders: An iterable of data_types.BuilderEntry containing the builders
- to query.
- builder_type: A string containing the type of builder to query, either
- "ci" or "try".
- Returns:
- A copy of |builders| with any inactive builders removed.
- """
- include_internal_builders = any(b.is_internal_builder for b in builders)
- query = self._GetActiveBuilderQuery(
- builder_type, include_internal_builders).encode('utf-8')
- cmd = GenerateBigQueryCommand(self._project, {}, batch=False)
- with open(os.devnull, 'w') as devnull:
- p = subprocess.Popen(cmd,
- stdout=subprocess.PIPE,
- stderr=devnull,
- stdin=subprocess.PIPE)
- stdout, _ = p.communicate(query)
- if not isinstance(stdout, six.string_types):
- stdout = stdout.decode('utf-8')
- results = json.loads(stdout)
- # We filter from an initial list instead of directly using the returned
- # builders since there are cases where they aren't equivalent, such as for
- # GPU tests if a particular builder doesn't run a particular suite. This
- # could be encapsulated in the query, but this would cause the query to take
- # longer. Since generating the initial list locally is basically
- # instantenous and we're optimizing for runtime, filtering is the better
- # option.
- active_builders = {r['builder_name'] for r in results}
- filtered_builders = [b for b in builders if b.name in active_builders]
- return filtered_builders
- def _QueryAddCombined(
- self,
- inputs: Tuple[data_types.BuilderEntry, data_types.TestExpectationMap]
- ) -> Tuple[data_types.ResultListType, str, data_types.TestExpectationMap]:
- """Combines the query and add steps for use in a process pool.
- Args:
- inputs: An iterable of inputs for QueryBuilder() and
- data_types.TestExpectationMap.AddResultList(). Should be in the order:
- builder expectation_map
- Returns:
- The output of data_types.TestExpectationMap.AddResultList().
- """
- builder, expectation_map = inputs
- results, expectation_files = self.QueryBuilder(builder)
- prefixed_builder_name = '%s/%s:%s' % (builder.project, builder.builder_type,
- builder.name)
- unmatched_results = expectation_map.AddResultList(prefixed_builder_name,
- results,
- expectation_files)
- return unmatched_results, prefixed_builder_name, expectation_map
- def QueryBuilder(self, builder: data_types.BuilderEntry
- ) -> Tuple[data_types.ResultListType, Optional[List[str]]]:
- """Queries ResultDB for results from |builder|.
- Args:
- builder: A data_types.BuilderEntry containing the builder to query.
- Returns:
- A tuple (results, expectation_files). |results| is the results returned by
- the query converted into a list of data_types.Result objects.
- |expectation_files| is a set of strings denoting which expectation files
- are relevant to |results|, or None if all should be used.
- """
- query_generator = self._GetQueryGeneratorForBuilder(builder)
- if not query_generator:
- # No affected tests on this builder, so early return.
- return [], None
- # Query for the test data from the builder, splitting the query if we run
- # into the BigQuery hard memory limit. Even if we keep failing, this will
- # eventually stop due to getting a QuerySplitError when we can't split the
- # query any further.
- query_results = None
- while query_results is None:
- try:
- query_results = self._RunBigQueryCommandsForJsonOutput(
- query_generator.GetQueries(), {
- '': {
- 'builder_name': builder.name
- },
- 'INT64': {
- 'num_builds': self._num_samples
- }
- })
- except MemoryLimitError:
- logging.warning(
- 'Query to builder %s hit BigQuery hard memory limit, trying again '
- 'with more query splitting.', builder.name)
- query_generator.SplitQuery()
- results = []
- if not query_results:
- # Don't bother logging if we know this is a fake CI builder.
- if not (builder.builder_type == constants.BuilderTypes.CI
- and builder in builders_module.GetInstance().GetFakeCiBuilders()):
- logging.warning(
- 'Did not get results for "%s", but this may be because its '
- 'results do not apply to any expectations for this suite.',
- builder.name)
- return results, None
- # It's possible that a builder runs multiple versions of a test with
- # different expectation files for each version. So, find a result for each
- # unique step and get the expectation files from all of them.
- results_for_each_step = {}
- for qr in query_results:
- step_name = qr['step_name']
- if step_name not in results_for_each_step:
- results_for_each_step[step_name] = qr
- expectation_files = []
- for qr in results_for_each_step.values():
- # None is a special value indicating "use all expectation files", so
- # handle that.
- ef = self._GetRelevantExpectationFilesForQueryResult(qr)
- if ef is None:
- expectation_files = None
- break
- expectation_files.extend(ef)
- if expectation_files is not None:
- expectation_files = list(set(expectation_files))
- for r in query_results:
- if self._ShouldSkipOverResult(r):
- continue
- results.append(self._ConvertJsonResultToResultObject(r))
- logging.debug('Got %d results for %s builder %s', len(results),
- builder.builder_type, builder.name)
- return results, expectation_files
- def _ConvertJsonResultToResultObject(self, json_result: QueryResult
- ) -> data_types.Result:
- """Converts a single BigQuery JSON result to a data_types.Result.
- Args:
- json_result: A single row/result from BigQuery in JSON format.
- Returns:
- A data_types.Result object containing the information from |json_result|.
- """
- build_id = _StripPrefixFromBuildId(json_result['id'])
- test_name = self._StripPrefixFromTestId(json_result['test_id'])
- actual_result = _ConvertActualResultToExpectationFileFormat(
- json_result['status'])
- tags = json_result['typ_tags']
- step = json_result['step_name']
- return data_types.Result(test_name, tags, actual_result, step, build_id)
- def _GetRelevantExpectationFilesForQueryResult(self, query_result: QueryResult
- ) -> Optional[Iterable[str]]:
- """Gets the relevant expectation file names for a given query result.
- Args:
- query_result: A dict containing single row/result from a BigQuery query.
- Returns:
- An iterable of strings containing expectation file names that are
- relevant to |query_result|, or None if all expectation files should be
- considered relevant.
- """
- raise NotImplementedError()
- def _ShouldSkipOverResult(self, result: QueryResult) -> bool:
- """Whether |result| should be ignored and skipped over.
- Args:
- result: A dict containing a single BigQuery result row.
- Returns:
- True if the result should be skipped over/ignored, otherwise False.
- """
- del result
- return False
- def _GetQueryGeneratorForBuilder(self, builder: data_types.BuilderEntry
- ) -> Optional['BaseQueryGenerator']:
- """Returns a BaseQueryGenerator instance to only include relevant tests.
- Args:
- builder: A data_types.BuilderEntry containing the builder to query.
- Returns:
- None if the query returned no results. Otherwise, some instance of a
- BaseQueryGenerator.
- """
- raise NotImplementedError()
- def _RunBigQueryCommandsForJsonOutput(self, queries: Union[str, List[str]],
- parameters: QueryParameters
- ) -> List[QueryResult]:
- """Runs the given BigQuery queries and returns their outputs as JSON.
- Args:
- queries: A string or list of strings containing valid BigQuery queries to
- run or a single string containing a query.
- parameters: A dict specifying parameters to substitute in the query in
- the format {type: {key: value}}. For example, the dict:
- {'INT64': {'num_builds': 5}}
- would result in --parameter=num_builds:INT64:5 being passed to
- BigQuery.
- Returns:
- The combined results of |queries| in JSON.
- """
- if isinstance(queries, str):
- queries = [queries]
- assert isinstance(queries, list)
- processes = set()
- processes_lock = threading.Lock()
- def run_cmd_in_thread(inputs: Tuple[List[str], str]) -> str:
- cmd, query = inputs
- query = query.encode('utf-8')
- with open(os.devnull, 'w') as devnull:
- with processes_lock:
- # Starting many queries at once causes us to hit rate limits much more
- # frequently, so stagger query starts to help avoid that.
- time.sleep(QUERY_DELAY)
- p = subprocess.Popen(cmd,
- stdout=subprocess.PIPE,
- stderr=devnull,
- stdin=subprocess.PIPE)
- processes.add(p)
- # We pass in the query via stdin instead of including it on the
- # commandline because we can run into command length issues in large
- # query mode.
- stdout, _ = p.communicate(query)
- if not isinstance(stdout, six.string_types):
- stdout = stdout.decode('utf-8')
- if p.returncode:
- # When running many queries in parallel, it's possible to hit the
- # rate limit for the account if we're unlucky, so try again if we do.
- if 'Exceeded rate limits' in stdout:
- raise RateLimitError()
- error_msg = 'Error running command %s. stdout: %s' % (cmd, stdout)
- if 'memory' in stdout:
- raise MemoryLimitError(error_msg)
- raise RuntimeError(error_msg)
- return stdout
- def run_cmd(cmd: List[str], tries: int) -> List[str]:
- if tries >= MAX_QUERY_TRIES:
- raise RuntimeError('Query failed too many times, aborting')
- # We use a thread pool with a thread for each query/process instead of
- # just creating the processes due to guidance from the Python docs:
- # https://docs.python.org/3/library/subprocess.html#subprocess.Popen.stderr
- # We need to write to stdin to pass the query in, but using
- # stdout/stderr/stdin directly is discouraged due to the potential for
- # deadlocks. The suggested method (using .communicate()) blocks, so we
- # need the thread pool to maintain parallelism.
- pool = multiprocessing.pool.ThreadPool(len(queries))
- def cleanup():
- pool.terminate()
- for p in processes:
- try:
- p.terminate()
- except OSError:
- # We can fail to terminate if the process is already finished, so
- # ignore such failures.
- pass
- processes.clear()
- args = [(cmd, q) for q in queries]
- try:
- return pool.map(run_cmd_in_thread, args)
- except RateLimitError:
- logging.warning('Query hit rate limit, retrying')
- cleanup()
- return run_cmd(cmd, tries + 1)
- finally:
- cleanup()
- raise RuntimeError('Hit branch that should be unreachable')
- bq_cmd = GenerateBigQueryCommand(self._project, parameters)
- stdouts = run_cmd(bq_cmd, 0)
- combined_json = []
- for result in [json.loads(s) for s in stdouts]:
- for row in result:
- combined_json.append(row)
- return combined_json
- def _StripPrefixFromTestId(self, test_id: str) -> str:
- """Strips the prefix from a test ID, leaving only the test case name.
- Args:
- test_id: A string containing a full ResultDB test ID, e.g.
- ninja://target/directory.suite.class.test_case
- Returns:
- A string containing the test cases name extracted from |test_id|.
- """
- raise NotImplementedError()
- def _GetActiveBuilderQuery(self, builder_type: str,
- include_internal_builders: bool) -> str:
- """Gets the SQL query for determining which builders actually produce data.
- Args:
- builder_type: A string containing the type of builders to query, either
- "ci" or "try".
- include_internal_builders: A boolean indicating whether internal builders
- should be included in the data that the query will access.
- Returns:
- A string containing a SQL query that will get all the names of all
- relevant builders that are active/producing data.
- """
- raise NotImplementedError()
- class BaseQueryGenerator(object):
- """Abstract base class for query generators."""
- def __init__(self, builder: data_types.BuilderEntry):
- self._builder = builder
- def SplitQuery(self) -> None:
- """Splits the query into more clauses/queries."""
- raise NotImplementedError('SplitQuery must be overridden in a child class')
- def GetClauses(self) -> List[str]:
- """Gets string representations of the test filters.
- Returns:
- A list of strings, each string being a valid SQL clause that applies a
- portion of the test filter to a query.
- """
- raise NotImplementedError('GetClauses must be overridden in a child class')
- def GetQueries(self) -> List[str]:
- """Gets string representations of the queries to run.
- Returns:
- A list of strings, each string being a valid SQL query that queries a
- portion of the tests of interest.
- """
- raise NotImplementedError('GetQueries must be overridden in a child class')
- # pylint: disable=abstract-method
- class FixedQueryGenerator(BaseQueryGenerator):
- """Concrete test filter that cannot be split."""
- def __init__(self, builder: data_types.BuilderEntry, test_filter: str):
- """
- Args:
- test_filter: A string containing the test filter SQL clause to use.
- """
- super(FixedQueryGenerator, self).__init__(builder)
- self._test_filter = test_filter
- def SplitQuery(self) -> None:
- raise QuerySplitError('Tried to split a query without any test IDs to use, '
- 'use --large-query-mode')
- def GetClauses(self) -> List[str]:
- return [self._test_filter]
- # pylint: enable=abstract-method
- # pylint: disable=abstract-method
- class SplitQueryGenerator(BaseQueryGenerator):
- """Concrete test filter that can be split to a desired size."""
- def __init__(self, builder: data_types.BuilderEntry, test_ids: List[str],
- target_num_samples: int):
- """
- Args:
- test_ids: A list of strings containing the test IDs to use in the test
- test filter.
- target_num_samples: The target/max number of samples to get from each
- query that uses clauses from this test filter.
- """
- super(SplitQueryGenerator, self).__init__(builder)
- self._test_id_lists = []
- self._target_num_samples = target_num_samples
- self._clauses = []
- self._PerformInitialSplit(test_ids)
- def _PerformInitialSplit(self, test_ids: List[str]) -> None:
- """Evenly splits |test_ids| into lists that are ~|_target_num_samples| long
- Only to be called from the constructor.
- Args:
- test_ids: A list of test IDs to split and assign to the _test_id_lists
- member.
- """
- assert isinstance(test_ids[0], six.string_types)
- num_lists = int(math.ceil(float(len(test_ids)) / self._target_num_samples))
- list_size = int(math.ceil(float(len(test_ids)) / num_lists))
- split_lists = []
- start = 0
- for _ in range(num_lists):
- end = min(len(test_ids), start + list_size)
- split_lists.append(test_ids[start:end])
- start = end
- self._test_id_lists = split_lists
- self._GenerateClauses()
- def _GenerateClauses(self) -> None:
- test_filter_clauses = []
- for id_list in self._test_id_lists:
- clause = 'AND test_id IN UNNEST([%s])' % ', '.join(id_list)
- test_filter_clauses.append(clause)
- self._clauses = test_filter_clauses
- def SplitQuery(self) -> None:
- def _SplitListInHalf(l: list) -> Tuple[list, list]:
- assert len(l) > 1
- front = l[:len(l) // 2]
- back = l[len(l) // 2:]
- return front, back
- tmp_test_id_lists = []
- for til in self._test_id_lists:
- if len(til) <= 1:
- raise QuerySplitError(
- 'Cannot split query any further, try lowering --num-samples')
- front, back = _SplitListInHalf(til)
- tmp_test_id_lists.append(front)
- tmp_test_id_lists.append(back)
- self._test_id_lists = tmp_test_id_lists
- self._GenerateClauses()
- def GetClauses(self) -> List[str]:
- return self._clauses
- # pylint: enable=abstract-method
- def GenerateBigQueryCommand(project: str,
- parameters: QueryParameters,
- batch: bool = True) -> List[str]:
- """Generate a BigQuery commandline.
- Does not contain the actual query, as that is passed in via stdin.
- Args:
- project: A string containing the billing project to use for BigQuery.
- parameters: A dict specifying parameters to substitute in the query in
- the format {type: {key: value}}. For example, the dict:
- {'INT64': {'num_builds': 5}}
- would result in --parameter=num_builds:INT64:5 being passed to BigQuery.
- batch: Whether to run the query in batch mode or not. Batching adds some
- random amount of overhead since it means the query has to wait for idle
- resources, but also allows for much better parallelism.
- Returns:
- A list containing the BigQuery commandline, suitable to be passed to a
- method from the subprocess module.
- """
- cmd = [
- 'bq',
- 'query',
- '--max_rows=%d' % MAX_ROWS,
- '--format=json',
- '--project_id=%s' % project,
- '--use_legacy_sql=false',
- ]
- if batch:
- cmd.append('--batch')
- for parameter_type, parameter_pairs in parameters.items():
- for k, v in parameter_pairs.items():
- cmd.append('--parameter=%s:%s:%s' % (k, parameter_type, v))
- return cmd
- def _StripPrefixFromBuildId(build_id: str) -> str:
- # Build IDs provided by ResultDB are prefixed with "build-"
- split_id = build_id.split('-')
- assert len(split_id) == 2
- return split_id[-1]
- def _ConvertActualResultToExpectationFileFormat(actual_result: str) -> str:
- # Web tests use ResultDB's ABORT value for both test timeouts and device
- # failures, but Abort is not defined in typ. So, map it to timeout now.
- if actual_result == 'ABORT':
- actual_result = json_results.ResultType.Timeout
- # The result reported to ResultDB is in the format PASS/FAIL, while the
- # expected results in an expectation file are in the format Pass/Failure.
- return expectations_parser.RESULT_TAGS[actual_result]
- class RateLimitError(Exception):
- """Exception raised when BigQuery hits a rate limit error."""
- class MemoryLimitError(Exception):
- """Exception raised when BigQuery hits its hard memory limit."""
- class QuerySplitError(Exception):
- """Exception raised when a query cannot be split any further."""
|