diff --git a/esbuild/gdc_elasticsearch.py b/esbuild/gdc_elasticsearch.py index 90a68ceb..73d4218b 100644 --- a/esbuild/gdc_elasticsearch.py +++ b/esbuild/gdc_elasticsearch.py @@ -19,6 +19,7 @@ from gdcdatamodel.models import File from progressbar import ProgressBar, Percentage, Bar, ETA from psqlgraph import PsqlGraphDriver +from esbuild.graph.common.cache import CachedGraph # TODO this could probably be bumped now that the number of bulk # threads in the config is higher, c.f. @@ -44,6 +45,7 @@ def shouldnt_delete(node): delete them. """ + if isinstance(node, File) and node.derived_files: return True else: @@ -55,8 +57,7 @@ class GDCElasticsearch(object): """ """ - def __init__(self, converter_class, es=None, - index_base="gdc_from_graph"): + def __init__(self, converter_class, es=None, index_base="gdc_from_graph"): """Walks the graph to produce elasticsearch json documents. :param es: An instance of Elasticsearch class @@ -75,14 +76,20 @@ def __init__(self, converter_class, es=None, os.environ.get("ES_PASSWORD", "")), timeout=9999) - self.graph = PsqlGraphDriver( + self.psqlgraph_args = ( os.environ["PG_HOST"], os.environ["PG_USER"], os.environ["PG_PASS"], os.environ["PG_NAME"], ) - self.converter = converter_class(self.graph) + self.graph = PsqlGraphDriver(self.psqlgraph_args) + + caching_options = converter_class.get_caching_options() + cache = CachedGraph(self.graph, caching_options) + converter = converter_class(cache) + + self.converter = converter def go(self, roll_alias=True): self.log.info("Caching database") @@ -90,7 +97,7 @@ def go(self, roll_alias=True): # that the cached database and which nodes get deleted is # consistent with self.graph.session_scope() as session: - self.converter.cache_database() + self.converter.cache.cache_database() self.log.info("Querying for old nodes to delete") to_delete = self.graph.nodes().sysan({"to_delete": True}).all() to_delete = [n for n in to_delete if not shouldnt_delete(n)] diff --git a/esbuild/graph/active/builder.py b/esbuild/graph/active/builder.py index e75ae3db..b8039739 100644 --- a/esbuild/graph/active/builder.py +++ b/esbuild/graph/active/builder.py @@ -214,7 +214,7 @@ def get_parent_with_category(self, node, category): if l['dst_type']._dictionary['category'] == category ] - return self.neighbors_labeled(node, labels) + return self.cache.neighbors_labeled(node.node_id, labels) def get_child_with_category(self, node, category): """returns iterable of neighors from inbound edges with category""" @@ -224,7 +224,7 @@ def get_child_with_category(self, node, category): if l['src_type']._dictionary['category'] == category ] - return self.neighbors_labeled(node, labels) + return self.cache.neighbors_labeled(node.node_id, labels) def add_file_analysis(self, node, doc): """Add the 'analysis' that produced the current file""" @@ -315,7 +315,8 @@ def get_read_group_qc_docs(self, read_group): """Returns a list of documents for Read Group QCs""" read_group_qc_docs = [] - rg_qcs = self.neighbors_labeled(read_group, 'read_group_qc') + rg_qcs = self.cache.neighbors_labeled( + read_group.node_id, 'read_group_qc') for read_group_qc in rg_qcs: read_group_qc_docs.append(self._get_base_doc(read_group_qc)) @@ -329,7 +330,7 @@ def get_file_read_groups(self, node): """ paths = self.file_to_read_group_paths.get(node.label, []) - return set(self.walk_paths(node, paths)) + return set(self.cache.walk_paths(node.node_id, paths)) def get_analysis_read_groups(self, node): """Given a analysis node, traverse up the tree to read_groups: @@ -356,7 +357,7 @@ def get_simple_file_doc(self, node): doc['data_format'] = self.get_data_format(node) - for dst in self.neighbors_labeled(node, 'data_subtype'): + for dst in self.cache.neighbors_labeled(node.node_id, 'data_subtype'): doc['data_type'] = dst['name'] return doc @@ -372,7 +373,8 @@ def get_file_associated_entities(self, node): entity for rg in self.get_file_read_groups(node) for entity in - self.neighbors_labeled(rg, self.possible_associated_entites) + self.cache.neighbors_labeled( + rg.node_id, self.possible_associated_entites) ] # Add entities with one step through a data_file @@ -380,7 +382,8 @@ def get_file_associated_entities(self, node): entity for parent in self.get_parent_with_category(node, 'data_file') for entity in - self.neighbors_labeled(parent, self.possible_associated_entites) + self.cache.neighbors_labeled( + parent.node_id, self.possible_associated_entites) ] # Copy number paths @@ -389,7 +392,8 @@ def get_file_associated_entities(self, node): list_product([['aliquot']], self.aliquot_to_copy_number_paths) ] entities += [ - entity for entity in self.walk_paths(node, cnv_paths) + entity for entity in self.cache.walk_paths( + node.node_id, cnv_paths) ] return list(set(entities)) diff --git a/esbuild/graph/common/builder.py b/esbuild/graph/common/builder.py index c32fe105..8d889af6 100644 --- a/esbuild/graph/common/builder.py +++ b/esbuild/graph/common/builder.py @@ -11,33 +11,150 @@ from cdisutils.log import get_logger from collections import defaultdict from copy import copy, deepcopy -from datadog import statsd +from functools32 import lru_cache from gdcdatamodel import models as md +from multiprocessing import cpu_count, Queue, Process from psqlgraph import Node, Edge from sqlalchemy.orm import joinedload -import itertools import logging -import networkx as nx import random import re +from esbuild.graph.common import ( + util, +) + +from esbuild.graph.common.index import ( + MemoryGraphIndex, + DiskGraphIndex, +) + +from esbuild.graph.common.cache import ( + CachingOptions, + CachedGraph, + CacheManager, +) + from .mappings import ( ONE_TO_MANY, ONE_TO_ONE, ) -from progressbar import ( - ProgressBar, - Percentage, - Bar, - ETA, -) - log = get_logger("graph_index") log.setLevel(level=logging.INFO) +def build_worker(builder, case_in_q, result_q): + """TODO: docstring + + """ + + while True: + + case = case_in_q.get() + if case is None: + return log.info('No more work for builder %s', builder) + + try: + result = builder.denormalize_case(case) + except Exception as exception: + log.exception(exception) + result = exception + + result_q.put(result) + del result + + +def start_worker_pool(builders, cases): + """Setup a process pool and schedule work to the case_in_q""" + + case_in_q, result_q = Queue(), Queue() + + pool = [ + Process( + target=build_worker, + args=(builder, case_in_q, result_q) + ) for builder in builders + ] + + # Schedule work + for case in cases: + case_in_q.put(case) + + # Put an end of work marker for all workers + for _ in range(len(builders)*2): + case_in_q.put(None) + + # Start all of the processes + for process in pool: + process.start() + + return case_in_q, result_q, pool + + +def build_index(builder_class, psqlgraph_driver_args, data_dir, cases=None, + threads=16): + """TODO: docstring + + """ + + index = DiskGraphIndex(data_dir) + + # Create managed cache + caching_options = builder_class.get_caching_options() + + cache = CachedGraph( + caching_options=caching_options, + psqlgraph_driver_args=psqlgraph_driver_args, + ) + + cache.cache_database() + + # Map work to worker processes + cases = cases or cache.get_cases() + builders = [builder_class(cache, index) for _ in range(threads)] + _, result_q, pool = start_worker_pool(builders, cases) + + pbar = util.get_pbar('Denormalizing cases ', len(cases)) + + # Collect results + while index.case_doc_count() < len(cases): + try: + if result_q.qsize() > 20: + log.warning("Primary thread overworked! %d", result_q.qsize()) + except NotImplementedError: + pass # on Mac OSX because of broken sem_getvalue() + + result = result_q.get() + if isinstance(result, Exception): + raise result + + case_doc, file_docs, annotation_docs = result + + # Collect docs + index.add_case_doc(case_doc) + map(index.add_annotation_doc, annotation_docs) + map(index.add_file_doc, file_docs) + + del result + del case_doc + del file_docs + del annotation_docs + + pbar.update(pbar.currval+1) + pbar.finish() + + # for process in pool: + # process.join() + + # Create project docs serially + project_docs = builders[0].denormalize_projects() + map(index.add_project_doc, project_docs) + + return index + + class GraphIndexBuilder(object): """This class handles all of the JSON production for the GDC @@ -78,18 +195,8 @@ class GraphIndexBuilder(object): remembering which nodes are walked through a lot and remembering which neighbors they have with a given label. - NOTE: An attempt was made to do this whole thing in parallel, - however the memory footprint grew to large. The best method for - doing this is to use the main process as a workload distributer, - and have child processes denormalizing cases. This way, - the main thread can upsert files on an outbound queue from child - processes. - - Josh (jsmiller@uchicago.edu) - TODOS: - - figure out a way to parallelize without excess copies - =============== Transformations =============== @@ -144,11 +251,9 @@ class GraphIndexBuilder(object): # "label": [{"key1": "value1", "key2": "value2"}] } - required_attrs = [ - 'mapper', - 'case_to_file_paths', - 'file_labels', - ] + # Suppress entities with redaction annotation if + # entity.annotation.category not in this list + redacted_but_not_suppressed = ['Subject withdrew consent'] supplement_regexes = [ re.compile(regex) for regex in [ @@ -161,12 +266,78 @@ class GraphIndexBuilder(object): ] ] - def __init__(self, psqlgraph_driver): + leaf_nodes = ['center', 'tissue_source_site'] + + # Omit entities from these projects + omitted_projects = { + ('TCGA', 'CNTL'), + ('TCGA', 'MISC'), + ('TCGA', 'TEST'), + ('TCGA', 'DEV1'), + ('TCGA', 'DEV2'), + ('TCGA', 'DEV3'), + ('TCGA', 'FPPP'), + ('GDC', 'INTERNAL'), + ('UAT08', 'BROAD-BCR'), + ('TARGET', 'AML-IF'), + } + + # The body of these nested documents will be flattened into + # the parent document using the given key's value + flatten = { + 'tag': 'name', + 'platform': 'name', + 'data_format': 'name', + 'data_subtype': 'name', + 'experimental_strategy': 'name', + 'data_level': 'name', + } + + # The edges below will maintain labels in the in memory graph, + # all others will be discarded + differentiated_edges = [ + ('file', 'member_of', 'archive'), + ('archive', 'member_of', 'file'), + ('file', 'describes', 'case'), + ('case', 'describes', 'file'), + ('file', 'related_to', 'file'), + ] + + possible_associated_entites = [ + 'portion', + 'aliquot', + 'case', + 'slide', + ] + + index_file_extensions = { + '.bai', + '.tbi', + } + + # The following attributes must be overridden + case_to_file_paths = None + mapper = None + file_labels = None + + required_attrs = [ + 'mapper', + 'case_to_file_paths', + 'file_labels', + ] + + def __init__(self, cached_graph, index): """Walks the graph to produce elasticsearch json documents. + :param cached_graph: Instance of CachedGraph (post .cached_database()) + :param index: Instance of GraphIndex + """ - # Verify required attributes are set + self.index = index + self.cache = cached_graph + + # verify required attributes are set for required_attr in self.required_attrs: if getattr(self, required_attr) is None: raise NotImplementedError( @@ -188,116 +359,36 @@ def __init__(self, psqlgraph_driver): # Get the actual case mapping to validate against self.case_es_mapping = self.mapper.get_case_es_mapping() - self.g = psqlgraph_driver - self.G = nx.Graph() - - self.leaf_nodes = ['center', 'tissue_source_site'] - self.experimental_strategies = {} - self.data_categories = {} - self.popular_nodes = {} - self.cases = None - self.projects = None - self.relevant_nodes = None - self.annotations = None - self.annotation_entities = None - self.entity_cases = None - - # Different from ``self.data_categories`` in that it's a - # replacement for a hardcoded dict of data_type, data_subtype - # relationships. This is populated by - # ``self._cache_existing_data_types()`` - self.existing_data_types = {} - - # Suppress entities with redaction annotation if - # entity.annotation.category not in this list - self.redacted_but_not_suppressed = ['Subject withdrew consent'] - - # Omit entities from these projects - self.omitted_projects = { - ('TCGA', 'CNTL'), - ('TCGA', 'MISC'), - ('TCGA', 'TEST'), - ('TCGA', 'DEV1'), - ('TCGA', 'DEV2'), - ('TCGA', 'DEV3'), - ('TCGA', 'FPPP'), - ('GDC', 'INTERNAL'), - ('UAT08', 'BROAD-BCR'), - ('TARGET', 'AML-IF'), - } - - # The body of these nested documents will be flattened into - # the parent document using the given key's value - self.flatten = { - 'tag': 'name', - 'platform': 'name', - 'data_format': 'name', - 'data_subtype': 'name', - 'experimental_strategy': 'name', - 'data_level': 'name', - } - - # The edges below will maintain labels in the in memory graph, - # all others will be discarded - self.differentiated_edges = [ - ('file', 'member_of', 'archive'), - ('archive', 'member_of', 'file'), - ('file', 'describes', 'case'), - ('case', 'describes', 'file'), - ('file', 'related_to', 'file'), - ] - - self.file_to_case_paths = [ - list(reversed(l))[1:]+['case'] - for l in self.case_to_file_paths - ] - - self.possible_associated_entites = [ - 'portion', - 'aliquot', - 'case', - 'slide', - ] - - self.index_file_extensions = { - '.bai', - '.tbi', - } - - def warning(self, title, text, tags=[], *args, **kwargs): - log.warning("{}: {}".format(title, text)) - statsd.event( - title, - text, - source_type_name="esbuild", - alert_type="warning", - tags=tags, + self.file_to_case_paths = util.reverse_paths( + self.case_to_file_paths, 'case') + + @classmethod + def get_caching_options(cls): + """Returns caching options for this builder""" + + return CachingOptions( + case_to_file_paths=cls.case_to_file_paths, + redacted_but_not_suppressed=cls.redacted_but_not_suppressed, + differentiated_edges=cls.differentiated_edges, + file_labels=cls.file_labels, + unindexed_by_property=cls.unindexed_by_property, + omitted_projects=cls.omitted_projects, + index_file_extensions=cls.index_file_extensions, + possible_associated_entites=cls.possible_associated_entites, + supplement_regexes=cls.supplement_regexes, ) - def error(self, title, text, tags=[], *args, **kwargs): - log.error("{}: {}".format(title, text)) - statsd.event( - title, - text, - source_type_name="esbuild", - alert_type="error", - tags=tags, - ) + @staticmethod + def warning(*args, **kwargs): + """Log a warning to logger and statsd""" - def pbar(self, title, maxval): - """Create and initialize a custom progressbar + util.log_warning(log, *args, **kwargs) - :param str title: The text of the progress bar - :param int maxval: The maximumum value of the progress bar + @staticmethod + def error(*args, **kwargs): + """Log a error to logger and statsd""" - """ - maxval = maxval or 1 # prevent maxal of 0 - pbar = ProgressBar(widgets=[ - title, Percentage(), ' ', - Bar(marker='#', left='[', right=']'), ' ', - ETA(), ' '], maxval=maxval) - pbar.update(0) - return pbar + util.log_error(log, *args, **kwargs) ################################################################### # Tree functions @@ -323,8 +414,8 @@ def create_tree(self, node, mapping, tree): if node.label in self.leaf_nodes: return {} submap = mapping[node.label] - corr, plural = submap['corr'] - for child in self.G.neighbors(node): + + for child in self.cache.neighbors(node.node_id): if child.label not in submap: continue tree[child] = {} @@ -338,7 +429,7 @@ def walk_tree(self, node, tree, mapping, doc, level=0, """ - corr, plural = mapping[node.label]['corr'] + corr, _ = mapping[node.label]['corr'] subdoc = self._get_base_doc(node) for child in tree[node]: child_corr, child_plural = mapping[node.label][child.label]['corr'] @@ -405,45 +496,7 @@ def _get_base_doc(self, node, include_id=True): return base ################################################################### - # Path functions - ################################################################### - - def walk_path(self, node, path, whole=False): - """Given a list of strings, treat it as a path, and yield the end of - possible traversals. If `whole` is true, return every node - along the traversal. - - """ - - if path: - for neighbor in self.neighbors_labeled(node, path[0]): - if whole or (len(path) == 1 and path[0] == neighbor.label): - yield neighbor - - for n in self.walk_path(neighbor, path[1:], whole): - yield n - - def walk_paths(self, node, paths, whole=False): - """Given a list of paths, yield the result of walking each path. If - `whole` is true, return every node along each traversal. - - """ - - return { - n for n in itertools.chain(*[ - self.walk_path(node, path, whole=whole) - for path in paths - ]) - } - - def remove_bam_index_files(self, files): - return { - f for f in files - if not self.is_index_file(f) - } - - ################################################################### - # Cases + # cases ################################################################## def remove_hidden_nodes(self, nodes): @@ -463,12 +516,17 @@ def remove_hidden_nodes(self, nodes): def get_case_files(self, node): """Return a list of file nodes by walking out from case""" - files = self.walk_paths(node, self.case_to_file_paths) - files = self.remove_bam_index_files(files) + files = self.cache.walk_paths(node.node_id, self.case_to_file_paths) + files = self.remove_index_files(files) files = self.remove_hidden_nodes(files) return files + def remove_index_files(self, files): + """Partial function for util.remove_index_files""" + + return util.remove_index_files(files, self.index_file_extensions) + def get_case_tree(self, node): """Use tree to create nested json @@ -611,8 +669,9 @@ def get_exp_strats(self, files): experimental_strategy is non-null """ - self._cache_experimental_strategies() - for exp_strat, file_list in self.experimental_strategies.iteritems(): + + exp_strats = self.cache.get_experimental_strategies() + for exp_strat, file_list in exp_strats.iteritems(): intersection = (file_list & files) if intersection: yield { @@ -626,8 +685,9 @@ def get_data_categories(self, files): data_category is non-null """ - self._cache_data_categories() - for data_category, file_list in self.data_categories.iteritems(): + data_categories = self.cache.get_data_categories() + + for data_category, file_list in data_categories.iteritems(): intersection = (file_list & files) if intersection: yield { @@ -747,7 +807,8 @@ def get_data_format(self, node): else: # get data_format from edge to DataFormat - formats = list(self.neighbors_labeled(node, 'data_format')) + formats = list(self.cache.neighbors_labeled( + node.node_id, 'data_format')) # Get the first format if formats: @@ -801,9 +862,14 @@ def add_file_neighbors(self, node, doc): """ - auto_neighbors = [n for n in dict(self.ftree_mapping['file']).keys() - if n not in ['archive', 'portion', 'file']] - for neighbor in set(self.neighbors_labeled(node, auto_neighbors)): + auto_neighbors = [ + n for n in dict(self.ftree_mapping['file']).keys() + if n not in ['archive', 'portion', 'file'] + ] + neighbors = set(self.cache.neighbors_labeled( + node.node_id, auto_neighbors)) + + for neighbor in neighbors: corr, label = self.ftree_mapping['file'][neighbor.label]['corr'] if neighbor.label in self.flatten: base = neighbor[self.flatten[neighbor.label]] @@ -831,24 +897,14 @@ def is_index_file(self, node): """ - # Active index files - if node._dictionary['category'] == 'index_file': - return True - - # Legacy index files - elif node.label == 'file': - for extension in self.index_file_extensions: - if node._props.get('file_name', '').endswith(extension): - return True - - else: - return False + return util.is_index_file(node, self.index_file_extensions) def get_file_index_files(self, node): """Given a file, return any neighboring index files""" return [ - n for n in list(self.neighbors_labeled(node, 'file')) - if self.G[node][n].get("label") == "related_to" + n for n in list(self.cache.neighbors_labeled(node.node_id, 'file')) + if self.cache.get_edge(node.node_id, n.node_id) + .get("label") == "related_to" and self.is_index_file(n) ] @@ -890,20 +946,23 @@ def add_related_files(self, node, doc): # Get related_files related_files = [ - n for n in list(self.neighbors_labeled(node, 'file')) - if self.G[node][n].get("label") == "related_to" + n for n in list(self.cache.neighbors_labeled( + node.node_id, 'file')) + if self.cache.get_edge(node.node_id, n.node_id) + .get("label") == "related_to" and not self.is_index_file(n) ] - related_files += list(self.neighbors_labeled(node, metadata_labels)) + related_files += list(self.cache.neighbors_labeled( + node.node_id, metadata_labels)) for related_file in related_files: rf_doc = self._get_base_doc(related_file, include_id=False) rf_doc['file_id'] = related_file.node_id # Data types - data_subtypes = self.neighbors_labeled( - related_file, + data_subtypes = self.cache.neighbors_labeled( + related_file.node_id, 'data_subtype', ) @@ -931,8 +990,10 @@ def add_related_files(self, node, doc): # file.archives) and one that is `related_to` (which goes # here). For now, we don't do this for non-legacy files. if node.label == 'file': - for archive in set(self.neighbors_labeled(node, 'archive')): - if self.G[node][archive].get('label') != 'member_of': + archives = set(self.cache.neighbors_labeled(node.node_id, 'archive')) + for archive in archives: + edge = self.cache.get_edge(node.node_id, archive.node_id) + if edge.get('label') != 'member_of': name = '{}.{}.0.tar.gz'.format( archive['submitter_id'], archive['revision']) rf_docs.append({ @@ -954,7 +1015,7 @@ def add_archives(self, node, doc): """ - for archive in set(self.neighbors_labeled(node, 'archive')): + for archive in set(self.cache.neighbors_labeled(node.node_id, 'archive')): if 'archive' in doc: return self.warning( "Duplicate archives for {}".format(node), @@ -964,7 +1025,8 @@ def add_archives(self, node, doc): is_skipped_legacy_edge = ( node.label == 'file' and - self.G[node][archive].get('label') != 'member_of' + self.cache.get_edge(node.node_id, archive) + .get('label') != 'member_of' ) if not is_skipped_legacy_edge: @@ -982,11 +1044,11 @@ def add_data_category(self, node, doc): """Add the data_subtype to the file document with child data_category """ + cached_data_categories = self.cache.get_data_categories() - self._cache_data_categories() data_categories = [ data_category - for data_category, files in self.data_categories.items() + for data_category, files in cached_data_categories.items() if node in files ] if data_categories: @@ -1005,11 +1067,11 @@ def add_cases(self, node, ptree, doc): log.warn('No ptree (case tree) for %s', node) return [] - if node not in self.relevant_nodes: + relevant = self.cache.get_relevant_nodes(node.node_id) + if not relevant: log.warn('No relevant cases for %s', node) return [] - relevant = self.relevant_nodes[node] prune_keys = ['sample', 'portion', 'analyte', 'aliquot', 'file'] self.prune_case(relevant, ptree, prune_keys) @@ -1025,6 +1087,19 @@ def add_cases(self, node, ptree, doc): return relevant + def get_node_annotation_docs(self, node_id): + """Returns the annotation docs for all annotations relevant to this + node + + """ + + annotations = self.cache.neighbors_labeled(node_id, 'annotation') + + return [ + self.denormalize_annotation(annotation) + for annotation in annotations + ] + def add_annotations(self, node, relevant, doc): """Given a file node, aggregate all of the annotations from a pruned case tree and insert them at the root level of the file @@ -1035,8 +1110,8 @@ def add_annotations(self, node, relevant, doc): annotations = doc.pop('annotations', []) for relevant_node in relevant: - ann_docs = self.annotation_entities.get(relevant_node, {}) - annotations.extend(ann_docs.values()) + annotations = self.get_node_annotation_docs(relevant_node.node_id) + annotations.extend(annotations) if annotations: doc['annotations'] = annotations @@ -1063,29 +1138,27 @@ def add_file_access(self, node, doc): def get_file_associated_entities(self, node): """Returns a list of entities that are 'associated' with a file""" - return list(self.neighbors_labeled( - node, self.possible_associated_entites)) + return list(self.cache.neighbors_labeled( + node.node_id, self.possible_associated_entites)) def add_file_associated_entities(self, node, doc, case_id): - self._cache_entity_cases() - docs = [] entities = self.get_file_associated_entities(node) - for e in entities: + for entity in entities: - if e not in self.entity_cases: + case = self.cache.get_entity_case(entity.node_id) + if not case: # Skip, the cases is likely missing because it is omitted continue - case = self.entity_cases[e] subdoc = { - 'entity_type': e.label, - 'entity_id': e.node_id, + 'entity_type': entity.label, + 'entity_id': entity.node_id, 'case_id': case.node_id } - entity_submitter_id = e._props.get('submitter_id') + entity_submitter_id = entity._props.get('submitter_id') if entity_submitter_id: subdoc['entity_submitter_id'] = entity_submitter_id @@ -1094,18 +1167,6 @@ def add_file_associated_entities(self, node, doc, case_id): if docs: doc['associated_entities'] = docs - def upsert_file_into_dict(self, files, file_doc): - did = file_doc['file_id'] - if did not in files: - files[did] = file_doc - else: - for case in file_doc['cases']: - case_id = case['case_id'] - existing_ids = { - p['case_id'] for p in files[did]['cases']} - if case_id not in existing_ids: - files[did]['cases'] += file_doc['cases'] - ################################################################### # Project summaries ################################################################### @@ -1114,11 +1175,11 @@ def denormalize_project(self, p): """Summarize a project. """ - self._cache_all() + doc = self._get_base_doc(p) # Get programs - program = self.neighbors_labeled(p, 'program').next() + program = self.cache.neighbors_labeled(p.node_id, 'program')[0] log.info('Program: {}'.format(program)) doc['program'] = self._get_base_doc(program) @@ -1126,7 +1187,7 @@ def denormalize_project(self, p): self.patch_project(doc) log.info('Finding cases') - cases = list(self.neighbors_labeled(p, 'case')) + cases = list(self.cache.neighbors_labeled(p.node_id, 'case')) log.info('Got {} cases'.format(len(cases))) # Get files @@ -1134,8 +1195,8 @@ def denormalize_project(self, p): files = set() case_files = {} for case in cases: - case_files[case] = self.remove_bam_index_files( - self.walk_paths(case, self.case_to_file_paths)) + case_files[case] = self.remove_index_files( + self.cache.walk_paths(case.node_id, self.case_to_file_paths)) files = files.union(case_files[case]) # filter files @@ -1148,11 +1209,11 @@ def denormalize_project(self, p): len(files), len(case_files))) # Get experimental strategies + experimental_strategies = self.cache.get_experimental_strategies() exp_strat_summaries = [] - self._cache_experimental_strategies() - for exp_strat in self.experimental_strategies.keys(): + for exp_strat in experimental_strategies: log.info('exp_strat: {}'.format(exp_strat)) - exp_files = (self.experimental_strategies[exp_strat] & files) + exp_files = (experimental_strategies[exp_strat] & files) if not len(exp_files): continue @@ -1170,11 +1231,11 @@ def denormalize_project(self, p): # Get data types data_category_summaries = [] - self._cache_data_categories() - for data_category in self.data_categories.keys(): + cached_data_categories = self.cache.get_data_categories() + for data_category in cached_data_categories.keys(): log.info('data_category: {}'.format(data_category)) - dt_files = (self.data_categories[data_category] & files) + dt_files = (cached_data_categories[data_category] & files) if not len(dt_files): continue @@ -1209,6 +1270,37 @@ def denormalize_project(self, p): return doc + def is_node_hidden(self, node): + """Return True if the node should be traversed (and therefore must + remain in the cache) but should not appear in any documents + + """ + + # Hide all submitted_* node types from indices + if node.label.startswith('submitted_'): + return True + + if node.label == 'archive': + return True + + return False + + @staticmethod + def node_labels_by_category(categories): + """Returns an iterator of node labels that are files + + """ + + categories = ( + tuple(categories) if hasattr(categories, '__iter__') + else (categories,) + ) + + return [ + n.label for n in Node.get_subclasses() + if n._dictionary['category'] in categories + ] + ################################################################### # Topmost denorm functions ################################################################### @@ -1223,19 +1315,21 @@ def denormalize_cases(self, cases=None): """ - self._cache_all() case_docs, ann_docs, file_docs = [], {}, {} if not cases: - cases = self.cases - pbar = self.pbar('Denormalizing cases ', len(cases)) + cases = self.cache.get_cases() + pbar = util.get_pbar('Denormalizing cases ', len(cases)) for n in cases: pa, fi, an = self.denormalize_case(n) case_docs.append(pa) + for a in an: if a['annotation_id'] not in ann_docs: ann_docs[a['annotation_id']] = a + for f in fi: - self.upsert_file_into_dict(file_docs, f) + util.upsert_file_into_dict(file_docs, f) + pbar.update(pbar.currval+1) pbar.finish() return case_docs, file_docs.values(), ann_docs.values() @@ -1247,11 +1341,11 @@ def denormalize_projects(self, projects=None): """ - self._cache_all() if not projects: - projects = self.projects + projects = self.cache.get_projects() + project_docs = [] - pbar = self.pbar('Denormalizing projects ', len(projects)) + pbar = util.get_pbar('Denormalizing projects ', len(projects)) for project in projects: project_docs.append(self.denormalize_project(project)) pbar.update(pbar.currval+1) @@ -1265,8 +1359,10 @@ def denormalize_annotation(self, node): case denormalization. """ + ann_doc = self._get_base_doc(node) - entities = self.G.neighbors(node) + entities = self.cache.neighbors(node.node_id) + if len(entities) == 0: self.error( 'Annotation has no entities', @@ -1303,93 +1399,16 @@ def denormalize_all(self): project documents """ - cases, files, annotations = self.denormalize_cases() - projects = self.denormalize_projects() - return cases, files, annotations, projects - - def denormalize_cases_sample(self, k=10): - """Return an entire index worth of case, file, annotation - documents - """ - self._cache_all() - cases = random.sample(self.cases, k) - cases, files, annotations = self.denormalize_cases(cases) - return cases, files, annotations + case_docs, file_docs, annotation_docs = self.denormalize_cases() + project_docs = self.denormalize_projects() - def denormalize_sample(self, k=10): - """Return an entire index worth of case, file, annotation, and - project documents + map(self.index.add_case_doc, case_docs) + map(self.index.add_file_doc, file_docs) + map(self.index.add_annotation_doc, annotation_docs) + map(self.index.add_project_doc, project_docs) - """ - cases, files, annotations = self.denormalize_sample_cases(k) - projs = random.sample(self.projects, 1) - projects = self.denormalize_projects(projs) - return cases, files, annotations, projects - - ################################################################### - # Graph functions - ################################################################### - - def nodes_labeled(self, labels): - """Returns an iterator over the edges in the graph with label `label` - - """ - - labels = tuple(labels) if hasattr(labels, '__iter__') else (labels,) - for n, p in self.G.nodes_iter(data=True): - if n.label in labels: - yield n - - @staticmethod - def node_labels_by_category(categories): - """Returns an iterator of node labels that are files - - """ - - categories = ( - tuple(categories) if hasattr(categories, '__iter__') - else (categories,) - ) - - return [ - n.label for n in Node.get_subclasses() - if n._dictionary['category'] in categories - ] - - def neighbors_labeled(self, node, labels, expected=None): - """For a given node, return an iterator with generates neighbors to - that node that are in a list of labels. `label` can be either a - string or list of strings. - - :param is_expected: Int count of expected elements - - """ - labels = tuple(labels) if hasattr(labels, '__iter__') else (labels,) - - if node in self.popular_nodes: - if labels not in self.popular_nodes[node]: - neighbors = self._cache_popular_neighbor( - node, self.G.neighbors(node), labels) - else: - neighbors = self.popular_nodes[node][labels] - else: - temp = self.G.neighbors(node) - if len(temp) > 200: - neighbors = self._cache_popular_neighbor(node, temp, labels) - else: - neighbors = {n for n in temp if n.label in labels} - - count = 0 - for n in neighbors: - count += 1 - yield n - - if expected is not None and count != expected: - self.warning( - "{}: unexpected no. of '{}' neighbors".format(node, labels), - '{}: {} != {} (expected)'.format(node, count, expected), - tags=["{}:{}".format(node.label, node.node_id)]) + return self.index ################################################################### # Validation functions @@ -1431,7 +1450,7 @@ def validate_annotations(self, ann_docs): ) def verify_data_category_count(self, case): - for data_category in self.existing_data_types.keys(): + for data_category in self.cache.get_existing_data_types(): calc = len([ f for f in case['files'] if f.get('data_category') == data_category @@ -1483,508 +1502,11 @@ def validate_case(self, node, case): if len(case['files']) != case['summary']['file_count']: self.error( 'Inconsistent case file count', - '{}: {} != {}'.format(node.node_id, len(case['files']), - case['summary']['file_count']), + '{}: {} != {}'.format( + node.node_id, len(case['files']), + case['summary']['file_count']), tags=["case_id:{}".format(node.node_id)], ) # Check for keys that are in the doc but not in the mapping self.validate_against_mapping(case, self.case_es_mapping) - - ################################################################### - # Caching functions - ################################################################### - - @staticmethod - def is_harmonized_file(node): - return ( - node.label == 'file' and - node._sysan.get('source', '').endswith('_alignment') - ) - - def is_old_supplement_file(self, node): - return ( - node.label == 'file' and - any( - p.match(node._props.get('file_name', '')) - for p in self.supplement_regexes - ) - ) - - def is_file_indexed(self, node): - """Returns false if node is a file that is not supposed to be indexed. - - """ - - # This function should only be for files - if node.label not in self.file_labels: - return True - - # Remove files with no acl entries - if len(node.acl) == 0: - log.info('File not indexed (empty acl): %s', node) - return False - - # Skip old versions of supplement xmls - if self.is_old_supplement_file(node): - log.info('File not indexed (deprecated supplement): %s', node) - return False - - # Skip old representation of harmonized files - if self.is_harmonized_file(node): - log.info('File not indexed (deprecated harmonized file): %s', node) - return False - - # Is file to_delete - if node.system_annotations.get("to_delete"): - return False - - # Is file not live - if node.state not in ['live', 'submitted']: - log.info('File not indexed (bad state: %s): %s', node, node.state) - return False - - return True - - def is_omitted_project_or_neighbor_case(self, node): - """Returns false if the node is a project that is not supposed to be - indexed. - - """ - - if node.label == 'project': - projects = [node] - elif node.label == 'case': - projects = list(self.neighbors_labeled(node, 'project', 1)) - else: - return False - - project_codes = [project.code for project in projects] - program_names = [ - program.name - for project in projects - for program in self.neighbors_labeled(project, 'program', 1) - ] - - # Check if project is not released - for project in projects: - if project.released is not True: - log.info('Omitting %s, project %s not released', node, project) - return True - - # Check project and program against omitted_projects - for program_name in program_names: - for project_code in project_codes: - if (program_name, project_code) in self.omitted_projects: - return True - - return False - - def is_unindexed_case(self, node): - return ( - node.label == 'case' - and not list(self.neighbors_labeled(node, 'project', 1)) - ) - - def is_node_unindexed_by_property(self, node): - """Returns True if node should be removed because its properties are - specified in self.unindexed_by_property as an indication to - remove it from the index. - - """ - - filters = self.unindexed_by_property.get(node.label, []) - - for filter_ in filters: - is_subset = not set(filter_.items()) - set(node._props.items()) - - if is_subset: - return True - - return False - - def is_node_indexed(self, node): - """Returns false if the node is not supposed to be indexed. - - """ - - if self.is_unindexed_case(node): - log.info('Node not indexed (case not indexed): {}'.format(node)) - return False - - # Check for non-indexed files - if not self.is_file_indexed(node): - log.info('Node not indexed (file not indexed): {}'.format(node)) - return False - - # Check for non-indexed files - if self.is_node_unindexed_by_property(node): - log.info('Node not indexed (not by property): {}'.format(node)) - return False - - # Check for omitted_projects - if self.is_omitted_project_or_neighbor_case(node): - log.info('Node not indexed (omitted project ): {}'.format(node)) - return False - - return True - - def is_node_hidden(self, node): - """Return True if the node should be traversed (and therefore must - remain in the cache) but should not appear in any documents - - """ - - # Hide all submitted_* node types from indices - if node.label.startswith('submitted_'): - return True - - if node.label == 'archive': - return True - - return False - - @staticmethod - def truncate_path(path, label): - """ - Given a path (a list of node labels), "truncate" it from the left - such that it starts with the given label, or return [], e.g.: - - truncate_path(["a", "b", "c"], "a") -> ["b", "c"] - truncate_path(["c", "d"], "b") -> [] - - """ - for i, currlabel in enumerate(path): - if currlabel == label: - return path[i+1:] - return [] - - def get_suppressed_children(self, redacted): - """Get the children of a redacted node. - - """ - to_suppress = [] - if redacted.label == "case": - paths = self.case_to_file_paths - else: - paths = [self.truncate_path(p, redacted.label) - for p in self.case_to_file_paths] - # filter empty paths - paths = [p for p in paths if p] - log.info("suppressing %s, which is redacted directly.", redacted) - to_suppress.append(redacted) - log.info("Walking down towards file with paths %s", paths) - extra = self.walk_paths(redacted, paths, whole=True) - log.info("Found %s other things to suppress by walking from %s", - extra, redacted) - to_suppress.extend(extra) - return to_suppress - - def suppressed_nodes(self): - """ - Find all nodes that need to be suppressed due to redactions. - """ - redactions = [a for a in self.nodes_labeled('annotation') - if a.classification == "Redaction" and - a.category not in self.redacted_but_not_suppressed] - to_suppress = [] - for redaction in redactions: - redacted_list = self.G.neighbors(redaction) - - if len(redacted_list) == 0: - # If there is no entity, then we have to move on to - # the next annotation - self.error( - 'Redaction annotation no entities', - "Redaction {} has zero entities associated.".format( - redaction), - tags=["annotation:{}".format(redaction)], - ) - continue - - if len(redacted_list) > 1: - # an annotation should only ever annotate one thing, - # however, proceed to redact them all - self.warning( - 'Redaction annotation has multiple entities', - ("{} has more than one entity associated. " - "For security reasons, removing all from index!") - .format(redaction), - tags=["annotation:{}".format(redaction)], - ) - - for redacted in redacted_list: - to_suppress += self.get_suppressed_children(redacted) - - # returning the redaction annotations themselves here might - # seem weird, but including the redaction annotations - # themselves without the things they point to won't work, so - # we have to remove them. - log.info("suppressing %s, the redaction annotation.", redaction) - to_suppress.append(redaction) - - return to_suppress - - def remove_unindexed_nodes_from_graph(self): - log.info('Selecting entities to be removed from cache...') - removed_nodes = [node for node in self.G.nodes() - if not self.is_node_indexed(node)] - log.info("Removing {} nodes from cache".format(len(removed_nodes))) - self.G.remove_nodes_from(removed_nodes) - log.info("Finding and removing suppressed nodes") - suppressed = self.suppressed_nodes() - log.info("Removing %s suppressed nodes", len(suppressed)) - self.G.remove_nodes_from(suppressed) - - def iter_database_edges(self): - """Returns an iterable of edges to load from the database. - - Eagerly (with join) loads the source and destination of the edge. - - """ - - return itertools.chain(*[ - self.g.edges(subclass) - .options(joinedload(subclass.src)) - .options(joinedload(subclass.dst)) - .yield_per(int(1e5)) - for subclass in Edge.__subclasses__() - ]) - - def cache_database(self): - """Load the database into memory and remember only edge labels that we - will need to distinguish later. - - """ - - with self.g.session_scope(): - pbar = self.pbar('Caching Database: ', self.g.edges().count()) - for e in self.iter_database_edges(): - pbar.update(pbar.currval+1) - triple = (e.src.label, e.label, e.dst.label) - needs_differentiation = (triple in self.differentiated_edges) - if triple == ("file", "data_from", "file"): - # for files that are "data_from" other files, the - # centers and aliquots of the source files count - # as neighbors of the dst files - for center in e.src.centers: - self.G.add_edge(e.dst, center) - for aliquot in e.src.aliquots: - self.G.add_edge(e.dst, aliquot) - if e.label == 'relates_to' and e.__dst_class__ == 'Case': - pass - elif needs_differentiation and e._props: - self.G.add_edge( - e.src, e.dst, label=e.label, props=e._props) - elif needs_differentiation and not e._props: - self.G.add_edge(e.src, e.dst, label=e.label) - elif e._props: - self.G.add_edge(e.src, e.dst, props=e._props) - else: - self.G.add_edge(e.src, e.dst) - pbar.finish() - - # Prune graph - log.info('Cached {} nodes'.format(self.G.number_of_nodes())) - self.remove_unindexed_nodes_from_graph() - - # Aggressively cache relationships, nodes by type, traversals, etc. - self._cache_all() - - def _cache_all(self): - """Create key value maps to cache nodes by label, by path, etc. - - """ - - self._cache_existing_data_types() - self._cache_experimental_strategies() - self._cache_data_categories() - self._cache_annotations() - self._cache_relevant_nodes() - self._cache_entity_cases() - self._cache_cases() - self._cache_projects() - - def _cache_projects(self): - """Save a list of all Project nodes""" - - if not self.projects: - log.info('Caching projects...') - self.projects = list(self.nodes_labeled('project')) - - def _cache_cases(self): - """Save a list of all Case nodes""" - - if not self.cases: - log.info('Caching cases...') - self.cases = list(self.nodes_labeled('case')) - - def _cache_entity_cases(self): - """Cache the related Case nodes for each file""" - - if self.entity_cases: - return - - entities = list(self.nodes_labeled(self.possible_associated_entites)) - pbar = self.pbar('Caching entity cases: ', len(entities)) - self.entity_cases = {} - - for e in entities: - if e.label == "case": - # if the associated entity is a case, it's case is - # just itself. this is kindy of sketchy but w/e - self.entity_cases[e] = e - continue - - paths = ( - self.truncate_path(path, e.label) - for path in self.file_to_case_paths - ) - cases = self.walk_paths(e, paths) - - if len(cases) > 1: - self.warning( - 'Entity associated with > 1 case', - '{}: Found {} cases'.format(e, len(cases)), - tags=["entity:{}".format(e)], - ) - return - - if len(cases) != 0: - self.entity_cases[e] = cases.pop() - - pbar.update(pbar.currval+1) - pbar.finish() - - def get_cls_file_to_case_paths(self, cls): - """Given a node, return the paths the lead monotonically up to case""" - - parent_labels = { - link['dst_type'].label - for link in cls._pg_links.values() - } - return ( - path - for path in self.file_to_case_paths - if path and path[0] in parent_labels - ) - - def _cache_relevant_nodes(self): - """The file documents will need to be pruned to only the nodes that - are relevant to the file. Here we cache all of the nodes - encountered when traversing to all related cases. - - """ - - if self.relevant_nodes: - return - - self.relevant_nodes = {} - - files = list(self.nodes_labeled(self.file_labels)) - pbar = self.pbar('Caching file paths: ', len(files)) - - for f in files: - paths = self.get_cls_file_to_case_paths(f) - self.relevant_nodes[f] = self.walk_paths(f, paths, whole=True) - pbar.update(pbar.currval+1) - - pbar.finish() - - def _cache_annotations(self): - if not self.annotations: - # cache what nodes are annotations - self.annotations = list(self.nodes_labeled('annotation')) - if self.annotation_entities: - # we've already cached the related entities - return - if not self.annotations: - # there aren't any entities to relate - self.annotation_entities = {} - log.warn('No annotations found in the cached database!') - return - pbar = self.pbar('Caching annotations: ', len(self.annotations)) - self.annotation_entities = {} - for a in self.annotations: - for n in self.G.neighbors(a): - if n not in self.annotation_entities: - self.annotation_entities[n] = {} - a_doc = self.denormalize_annotation(a) - self.annotation_entities[n][a.node_id] = a_doc - pbar.update(pbar.currval+1) - pbar.finish() - - def _cache_popular_neighbor(self, node, neighbors, labels): - if node not in self.popular_nodes: - self.popular_nodes[node] = {} - self.popular_nodes[node][labels] = { - n for n in neighbors if n.label in labels} - return self.popular_nodes[node][labels] - - def _cache_data_categories(self): - """Looking up the files that are classified in each data_type is a - common computation. Here we cache this information for easy retrieval. - - ..note:: - data_type is renamed data_category, viz. - https://jira.opensciencedatacloud.org/browse/PGDC-1472 - - """ - - if len(self.data_categories): - return - - log.info('Caching data categories') - for data_category in self.nodes_labeled('data_type'): - category = data_category._props['name'] - self.data_categories[category] = self.remove_bam_index_files( - set(self.walk_path(data_category, ['data_subtype', 'file']))) - - # New files have 'data_category' as a property - for file_ in self.nodes_labeled(self.file_labels): - category = file_._props.get('data_category') - if not category: - continue - self.data_categories.setdefault(category, set()).add(file_) - - def _cache_experimental_strategies(self): - """Looking up the files that are classified in each - experimental_strategy is a common computation. Here we cache - this information for easy retrieval. - - """ - - if len(self.experimental_strategies): - return - - log.info('Caching experitmental strategies') - for exp_strat in self.nodes_labeled('experimental_strategy'): - strategy = exp_strat._props['name'] - self.experimental_strategies[strategy] = set(self.walk_path( - exp_strat, ['file'])) - - # New files have 'experimental_strategy' as a property - for file_ in self.nodes_labeled(self.file_labels): - strategy = file_._props.get('experimental_strategy') - if not strategy: - continue - self.experimental_strategies.setdefault(strategy, set()).add(file_) - - def _cache_existing_data_types(self): - """The last version of this code imported a hard coded list and called - it DATA_TYPES. This function replaces this hardcoded nested - dict by pulling it from the graph at runtime. - - :returns: - The data types in the graph in the format - ``{'data_type.name': ['data_subtype.name']}`` - - """ - - with self.g.session_scope(): - return { - data_type.name: [ - subtype.name - for subtype in data_type.data_subtypes - ] for data_type in self.g.nodes(md.DataType).all() - } diff --git a/esbuild/graph/common/cache.py b/esbuild/graph/common/cache.py new file mode 100644 index 00000000..1c6a7519 --- /dev/null +++ b/esbuild/graph/common/cache.py @@ -0,0 +1,804 @@ +# -*- coding: utf-8 -*- +""" +esbuild.graph.common.cache +---------------------------------- + +Functionality to create a multiprocessing manager to share cached +graph information between processes + +""" + +import cdisutils +import gc +import itertools +import logging +import networkx as nx +import types + +from collections import namedtuple +from gdcdatamodel import models as md +from multiprocessing.managers import BaseManager +from multiprocessing import Pool, Queue, TimeoutError +from psqlgraph import Edge, Node, PsqlGraphDriver +from sqlalchemy.orm import joinedload +from types import StringTypes + +from esbuild.graph.common import ( + util, +) + +from esbuild.graph.common.fake_node import ( + FakeNode, +) + +logger = cdisutils.log.get_logger(__name__) +logger.setLevel(logging.INFO) + + +def is_case_cache_edge(edge): + """Determine if this edge or edge class is just a case cache edge""" + + return edge.label == 'relates_to' and edge.__dst_class__ == 'Case' + + +class CachingOptions(object): + + """An object to hold the options required to hold the required options + to cache information to a SharedGraph object + + """ + + def __init__( + self, + case_to_file_paths, + redacted_but_not_suppressed, + differentiated_edges, + file_labels, + unindexed_by_property, + omitted_projects, + index_file_extensions, + possible_associated_entites, + supplement_regexes, + ): + """Options required to cache information to SharedGraph object + + :param list case_to_file_paths: + specifies all possible paths from case entities to file + entities. is a list of lists of string labels + + :param list redacted_but_not_suppressed: + Suppress entities with redaction annotation if + entity.annotation.category not in this list + + :param file_labels: + Node labels that should be treated as properties + + :param unindezed_by_property: + Filter nodes out if their properties are a superset of any of + the dictionaries listed here by label + + """ + + self.case_to_file_paths = case_to_file_paths + self.redacted_but_not_suppressed = redacted_but_not_suppressed + self.differentiated_edges = differentiated_edges + self.file_labels = file_labels + self.unindexed_by_property = unindexed_by_property + self.omitted_projects = omitted_projects + self.index_file_extensions = index_file_extensions + self.possible_associated_entites = possible_associated_entites + self.supplement_regexes = supplement_regexes + + self.file_to_case_paths = util.reverse_paths( + self.case_to_file_paths, 'case') + + +class CachedGraph(object): + """Represents the shared information for esbuild workers. This data + includes a cached version of the psqlgraph graph in NetworkX as + well as more aggresively cached things, e.g. relationships. + + """ + + def __init__(self, caching_options, psqlgraph_driver_args=None, + psqlgraph_driver_kwargs=None, psqlgraph_driver=None): + + # Injected Dependencies + self.psqlgraph_driver = psqlgraph_driver or PsqlGraphDriver( + *(psqlgraph_driver_args or []), + **(psqlgraph_driver_kwargs or {}) + ) + + self.graph = nx.Graph() + self.caching_options = caching_options + + # Cached information + self.nodes = {} + self.experimental_strategies = {} + self.data_categories = {} + self.popular_nodes = {} + self.cases = None + self.projects = None + self.relevant_nodes = None + self.annotations = None + self.entity_cases = None + # Different from ``self.data_categories`` in that it's a + # replacement for a hardcoded dict of data_type, data_subtype + # relationships. This is populated by + # ``self._cache_existing_data_types()`` + self.existing_data_types = {} + + ################################################################### + # Path functions + ################################################################### + + def nodes_labeled(self, labels): + """Returns an iterator over the edges in the graph with label `label` + + """ + + labels = tuple(labels) if hasattr(labels, '__iter__') else (labels,) + for node, _ in self.graph.nodes_iter(data=True): + if node.label in labels: + yield node + + def neighbors_labeled(self, *args, **kwargs): + """Proxy for self._neighbors_labeled which was original written as a + generator to return a list instead + + """ + + return list(self._neighbors_labeled(*args, **kwargs)) + + def _neighbors_labeled(self, node_id, labels, expected=None): + + """For a given node, return an iterator with generates neighbors to + that node that are in a list of labels. `label` can be either a + string or list of strings. + + :param is_expected: Int count of expected elements + + """ + + labels = tuple(labels) if hasattr(labels, '__iter__') else (labels,) + node = self.get_node_in_graph(node_id) + + if node in self.popular_nodes: + if labels not in self.popular_nodes[node]: + neighbors = self._cache_popular_neighbor( + node, self.graph.neighbors(node), labels) + else: + neighbors = self.popular_nodes[node][labels] + + else: + temp = self.graph.neighbors(node) + if len(temp) > 200: + neighbors = self._cache_popular_neighbor(node, temp, labels) + else: + neighbors = {n for n in temp if n.label in labels} + + count = 0 + for neighbor in neighbors: + count += 1 + yield neighbor + + if expected is not None and count != expected: + self.warning( + "{}: unexpected no. of '{}' neighbors".format(node, labels), + '{}: {} != {} (expected)'.format(node, count, expected), + tags=["{}:{}".format(node.label, node.node_id)]) + + def neighbors(self, node_id): + """Return the neighbors of given node""" + + node = self.get_node_in_graph(node_id) + return self.graph.neighbors(node) + + def walk_path(self, node_id, path, whole=False): + """Given a list of strings, treat it as a path, and yield the end of + possible traversals. If `whole` is true, return every node + along the traversal. + + """ + node = self.get_node_in_graph(node_id) + + if path: + for neighbor in self.neighbors_labeled(node.node_id, path[0]): + if whole or (len(path) == 1 and path[0] == neighbor.label): + yield neighbor + + for node in self.walk_path(neighbor.node_id, path[1:], whole): + yield node + + def walk_paths(self, node_id, paths, whole=False): + """Given a list of paths, yield the result of walking each path. If + `whole` is true, return every node along each traversal. + + """ + + return { + n for n in itertools.chain(*[ + self.walk_path(node_id, path, whole=whole) + for path in paths if path + ]) + } + + ################################################################### + # Proxy Methods + ################################################################### + + def get_relevant_nodes(self, node_id): + """Return the nodes relevant to this one""" + + node = self.get_node_in_graph(node_id) + return self.relevant_nodes.get(node, []) + + def get_entity_case(self, node_id): + """Return the case associated with this node""" + + node = self.get_node_in_graph(node_id) + return self.entity_cases.get(node, None) + + @staticmethod + def warning(*args, **kwargs): + """Log a warning to logger and statsd""" + + util.log_warning(logger, *args, **kwargs) + + @staticmethod + def error(*args, **kwargs): + """Log a error to logger and statsd""" + + util.log_error(logger, *args, **kwargs) + + def get_cases(self): + """Proxy to get cached cases""" + + return self.cases + + def get_projects(self): + """Proxy to get cached projects""" + + return self.projects + + def get_experimental_strategies(self): + """Proxy to get cached experimental_strategies""" + + return self.experimental_strategies + + def get_existing_data_types(self): + """Proxy to get cached existing_data_types""" + + return self.existing_data_types + + def get_data_categories(self): + """Proxy to get cached data_categories""" + + return self.data_categories + + def get_edge(self, src_id, dst_id): + """Returns any information stored about the edge between two nodes""" + + src = self.get_node_in_graph(src_id) + dst = self.get_node_in_graph(dst_id) + + return self.graph[src][dst] + + def get_node_in_graph(self, node_id): + """Given a node or a node_id, return the corresponding node that is in + the NetworkX graph + + """ + + assert isinstance(node_id, types.StringTypes) + + return self.nodes[node_id] + + ################################################################### + # Setup Methods + ################################################################### + + def _get_suppressed_children(self, redacted): + """Get the children of a redacted node""" + + to_suppress = [] + if redacted.label == "case": + paths = self.caching_options.case_to_file_paths + else: + paths = [ + util.truncate_path(p, redacted.label) + for p in self.caching_options.case_to_file_paths if p + ] + + logger.info("suppressing %s, which is redacted directly.", redacted) + to_suppress.append(redacted) + + logger.info("Walking down towards file with paths %s", paths) + extra = self.walk_paths(redacted.node_id, paths, whole=True) + + logger.info("Found %s other things to suppress by walking from %s", + extra, redacted) + to_suppress.extend(extra) + + return to_suppress + + def _suppressed_nodes(self): + """ + Find all nodes that need to be suppressed due to redactions. + """ + + redactions = [ + annotation for annotation in self.nodes_labeled('annotation') + if annotation.classification == "Redaction" and + annotation.category not in + self.caching_options.redacted_but_not_suppressed + ] + + to_suppress = [] + + for redaction in redactions: + redacted_list = self.graph.neighbors(redaction) + + if len(redacted_list) == 0: + # If there is no entity, then we have to move on to + # the next annotation + self.error( + 'Redaction annotation no entities', + "Redaction {} has zero entities associated." + .format(redaction), + tags=["annotation:{}".format(redaction)], + ) + continue + + if len(redacted_list) > 1: + # an annotation should only ever annotate one thing, + # however, proceed to redact them all + self.warning( + 'Redaction annotation has multiple entities', + ("{} has more than one entity associated. " + "For security reasons, removing all from index!") + .format(redaction), + tags=["annotation:{}".format(redaction)], + ) + + for redacted in redacted_list: + to_suppress += self._get_suppressed_children(redacted) + + # returning the redaction annotations themselves here might + # seem weird, but including the redaction annotations + # themselves without the things they point to won't work, so + # we have to remove them. + logger.info("suppressing %s, the redaction annotation.", redaction) + to_suppress.append(redaction) + + return to_suppress + + def _is_unindexed_case(self, node): + return ( + node.label == 'case' + and not list(self.neighbors_labeled(node.node_id, 'project', 1)) + ) + + def _is_node_indexed(self, node): + """Returns false if the node is not supposed to be indexed. + + """ + + if self._is_unindexed_case(node): + logger.info('Node not indexed (case not indexed): %s', node) + return False + + # Check for non-indexed files + if not self.is_file_indexed(node): + logger.info('Node not indexed (file not indexed): %s', node) + return False + + # Check for non-indexed files + if self.is_node_unindexed_by_property(node): + logger.info('Node not indexed (not by property): %s', node) + return False + + # Check for omitted_projects + if self.is_omitted_project_or_neighbor_case(node): + logger.info('Node not indexed (omitted project ): %s', node) + return False + + return True + + def _remove_unindexed_nodes_from_graph(self): + logger.info('Selecting entities to be removed from cache...') + + removed_nodes = [ + node for node in self.graph.nodes() + if not self._is_node_indexed(node) + ] + + logger.info("Removing %s nodes from cache", len(removed_nodes)) + self.graph.remove_nodes_from(removed_nodes) + for node in removed_nodes: + self.nodes.pop(node.node_id, None) + + logger.info("Finding and removing suppressed nodes") + suppressed = self._suppressed_nodes() + + logger.info("Removing %s suppressed nodes", len(suppressed)) + self.graph.remove_nodes_from(suppressed) + for node in suppressed: + self.nodes.pop(node.node_id, None) + + def _iter_database_edges(self): + """Returns an iterable of edges to load from the database. + + Eagerly (with join) loads the source and destination of the edge. + + """ + + return itertools.chain(*[ + self.psqlgraph_driver.edges(subclass) + .options(joinedload(subclass.src)) + .options(joinedload(subclass.dst)) + .yield_per(int(1e5)) + for subclass in sorted(Edge.__subclasses__()) + if not is_case_cache_edge(subclass) + ]) + + + def _get_fake_node(self, node): + """If we've seen this node before, then return the FakeNode version of + it, otherwise create a new one and return that. + + """ + + existing = self.nodes.get(node.node_id) + if not existing: + existing = self.nodes.setdefault(node.node_id, FakeNode(node)) + + return existing + + def cache_database(self): + """Load the database into memory and remember only edge labels that we + will need to distinguish later. + + """ + + with self.psqlgraph_driver.session_scope() as session: + # Meter the progress bar by nodes, because creating the + # nodes will be the majority of the time + node_count = self.psqlgraph_driver.nodes().count() + pbar = util.get_pbar('Caching Database: ', node_count) + + for edge in self._iter_database_edges(): + pbar.update(len(self.nodes)) + + src = self._get_fake_node(edge.src) + dst = self._get_fake_node(edge.dst) + + triple = (src.label, edge.label, dst.label) + needs_differentiation = ( + triple in self.caching_options.differentiated_edges + ) + + if triple == ("file", "data_from", "file"): + # for files that are "data_from" other files, the + # centers and aliquots of the source files count + # as neighbors of the dst files + for center in edge.src.centers: + self.graph.add_edge(dst, self._get_fake_node(center)) + + for aliquot in edge.src.aliquots: + self.graph.add_edge(dst, self._get_fake_node(aliquot)) + + elif needs_differentiation and edge._props: + self.graph.add_edge( + src, dst, label=edge.label, props=edge._props) + + elif needs_differentiation and not edge._props: + self.graph.add_edge(src, dst, label=edge.label) + + elif edge._props: + self.graph.add_edge(src, dst, props=edge._props) + + else: + self.graph.add_edge(src, dst) + + session.expunge_all() + pbar.finish() + + # Prune graph + logger.info('Cached {} nodes'.format(self.graph.number_of_nodes())) + self._remove_unindexed_nodes_from_graph() + + # Aggressively cache relationships, nodes by type, traversals, etc. + self._cache_all() + + def _cache_all(self): + """Create key value maps to cache nodes by label, by path, etc. + + """ + + self._cache_node_ids() + self._cache_existing_data_types() + self._cache_experimental_strategies() + self._cache_data_categories() + self._cache_relevant_nodes() + self._cache_entity_cases() + self._cache_cases() + self._cache_projects() + + def _cache_node_ids(self): + """Create a hashtable from node_id to node in the graph""" + + self.nodes = { + node.node_id: node + for node in self.graph.nodes_iter() + } + + def _cache_projects(self): + """Save a list of all Project nodes""" + + if not self.projects: + logger.info('Caching projects...') + self.projects = list(self.nodes_labeled('project')) + + def _cache_cases(self): + """Save a list of all Case nodes""" + + if not self.cases: + logger.info('Caching cases...') + self.cases = list(self.nodes_labeled('case')) + + def _cache_entity_cases(self): + """Cache the related Case nodes for each file""" + + if self.entity_cases: + return + + entities = list(self.nodes_labeled( + self.caching_options.possible_associated_entites)) + pbar = util.get_pbar('Caching entity cases: ', len(entities)) + self.entity_cases = {} + + for entity in entities: + if entity.label == "case": + # if the associated entity is a case, it's case is + # just itself. this is kindy of sketchy but w/e + self.entity_cases[entity] = entity + continue + + paths = ( + util.truncate_path(path, entity.label) + for path in self.caching_options.file_to_case_paths + ) + cases = self.walk_paths(entity.node_id, paths) + + if len(cases) > 1: + self.warning( + 'Entity associated with > 1 case', + '{}: Found {} cases'.format(entity, len(cases)), + tags=["entity:{}".format(entity)], + ) + return + + if len(cases) != 0: + self.entity_cases[entity] = cases.pop() + + pbar.update(pbar.currval+1) + pbar.finish() + + def _cache_relevant_nodes(self): + """The file documents will need to be pruned to only the nodes that + are relevant to the file. Here we cache all of the nodes + encountered when traversing to all related cases. + + """ + + if self.relevant_nodes: + return + + self.relevant_nodes = {} + + files = list(self.nodes_labeled(self.caching_options.file_labels)) + pbar = util.get_pbar('Caching file paths: ', len(files)) + + for file_ in files: + paths = util.get_file_to_case_paths( + file_, self.caching_options.file_to_case_paths) + self.relevant_nodes[file_] = self.walk_paths( + file_.node_id, paths, whole=True) + pbar.update(pbar.currval+1) + + pbar.finish() + + def _cache_popular_neighbor(self, node, neighbors, labels): + if node not in self.popular_nodes: + self.popular_nodes[node] = {} + + self.popular_nodes[node][labels] = { + n for n in neighbors if n.label in labels + } + + return self.popular_nodes[node][labels] + + def _cache_data_categories(self): + """Looking up the files that are classified in each data_type is a + common computation. Here we cache this information for easy retrieval. + + ..note:: + data_type is renamed data_category, viz. + https://jira.opensciencedatacloud.org/browse/PGDC-1472 + + """ + + if len(self.data_categories): + return + + logger.info('Caching data categories') + for data_category in self.nodes_labeled('data_type'): + category = data_category._props['name'] + + self.data_categories[category] = util.remove_index_files( + set(self.walk_path( + data_category.node_id, ['data_subtype', 'file'])), + self.caching_options.index_file_extensions, + ) + + # New files have 'data_category' as a property + for file_ in self.nodes_labeled(self.caching_options.file_labels): + category = file_._props.get('data_category') + if not category: + continue + self.data_categories.setdefault(category, set()).add(file_) + def _cache_experimental_strategies(self): + """Looking up the files that are classified in each + experimental_strategy is a common computation. Here we cache + this information for easy retrieval. + + """ + + if len(self.experimental_strategies): + return + + logger.info('Caching experitmental strategies') + for exp_strat in self.nodes_labeled('experimental_strategy'): + strategy = exp_strat._props['name'] + self.experimental_strategies[strategy] = set(self.walk_path( + exp_strat.node_id, ['file'])) + + # New files have 'experimental_strategy' as a property + for file_ in self.nodes_labeled(self.caching_options.file_labels): + strategy = file_._props.get('experimental_strategy') + if not strategy: + continue + self.experimental_strategies.setdefault(strategy, set()).add(file_) + + def _cache_existing_data_types(self): + """The last version of this code imported a hard coded list and called + it DATA_TYPES. This function replaces this hardcoded nested + dict by pulling it from the graph at runtime. + + :returns: + The data types in the graph in the format + ``{'data_type.name': ['data_subtype.name']}`` + + """ + + with self.psqlgraph_driver.session_scope(): + return { + data_type.name: [ + subtype.name + for subtype in data_type.data_subtypes + ] for data_type in self.psqlgraph_driver.nodes(md.DataType) + } + + @staticmethod + def is_harmonized_file(node): + return ( + node.label == 'file' and + node._sysan.get('source', '').endswith('_alignment') + ) + + def is_old_supplement_file(self, node): + return ( + node.label == 'file' + and any( + pattern.match(node._props.get('file_name', '')) + for pattern in self.caching_options.supplement_regexes + ) + ) + + def is_file_indexed(self, node): + """Returns false if node is a file that is not supposed to be indexed. + + """ + + # This function should only be for files + if node.label not in self.caching_options.file_labels: + return True + + # Remove files with no acl entries + if len(node.acl) == 0: + logger.info('File not indexed (empty acl): %s', node) + return False + + # Skip old versions of supplement xmls + if self.is_old_supplement_file(node): + logger.info('File not indexed (deprecated supplement): %s', node) + return False + + # Skip old representation of harmonized files + if self.is_harmonized_file(node): + logger.info('File not indexed (deprecated harmonized file): %s', node) + return False + + # Is file to_delete + if node.system_annotations.get("to_delete"): + return False + + # Is file not live + if node.state not in ['live', 'submitted']: + logger.info('File not indexed (bad state: %s): %s', node, node.state) + return False + + return True + + def is_omitted_project_or_neighbor_case(self, node): + """Returns false if the node is a project that is not supposed to be + indexed. + + """ + + if node.label == 'project': + projects = [node] + elif node.label == 'case': + projects = self.neighbors_labeled(node.node_id, 'project', 1) + else: + return False + + project_codes = [project.code for project in projects] + program_names = [ + program.name + for project in projects + for program in self.neighbors_labeled(project.node_id, 'program', 1) + ] + + # Check if project is not released + for project in projects: + if project.released is not True: + logger.info('Omitting %s, project %s not released', + node, project) + return True + + # Check project and program against omitted_projects + for program_name in program_names: + for project_code in project_codes: + is_omitted = ( + (program_name, project_code) in + self.caching_options.omitted_projects + ) + if is_omitted: + return True + + return False + + def is_node_unindexed_by_property(self, node): + """Returns True if node should be removed because its properties are + specified in self.unindexed_by_property as an indication to + remove it from the index. + + """ + + filters = self.caching_options.unindexed_by_property.get(node.label, []) + + for filter_ in filters: + is_subset = not set(filter_.items()) - set(node._props.items()) + + if is_subset: + return True + + return False + +class CacheManager(BaseManager): + pass + + +CacheManager.register('new_cached_graph', CachedGraph) diff --git a/esbuild/graph/common/fake_node.py b/esbuild/graph/common/fake_node.py new file mode 100644 index 00000000..02b666e8 --- /dev/null +++ b/esbuild/graph/common/fake_node.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +""" +esbuild.graph.common.fake_node +---------------------------------- + +Create a smaller, faster, serializable psqlgraph.Node replacement + +""" + + +class FakeNode(object): + + def __init__(self, node): + props = dict(node.props) + sysan = dict(node._sysan) + acl = list(node.acl) + label = str(node.label) + node_id = str(node.node_id) + + self._class_name = str(node.__class__.__name__) + self.node_id = node_id + self.acl = acl + self.label = label + + self._props = props + self.properties = props + self.props = props + + self._sysan = sysan + self.system_annotations = sysan + self.sysan = sysan + + self.__pg_properties__ = node.__pg_properties__ + self._pg_links = node._pg_links + self._pg_backrefs = node._pg_backrefs + self._pg_edges = node._pg_edges + self._dictionary = node._dictionary + + for key, value in props.iteritems(): + setattr(self, key, value) + + def __repr__(self): + return '<{}({})>'.format(self._class_name, self.node_id) + + def __getitem__(self, key): + return self.props[key] + + def __eq__(self, other): + return bool(self.node_id == other.node_id) + + def __ne__(self, other): + return bool(self.node_id != other.node_id) diff --git a/esbuild/graph/common/index.py b/esbuild/graph/common/index.py new file mode 100644 index 00000000..3fff23ca --- /dev/null +++ b/esbuild/graph/common/index.py @@ -0,0 +1,299 @@ +# -*- coding: utf-8 -*- +""" +esbuild.graph.common.index +---------------------------------- + +Defines :class:`GraphIndex` represents a GDC Data Portal index + +""" + +import abc +import cdisutils +import json +import logging +import os +import shutil +import time + + +logger = cdisutils.log.get_logger(__name__) +logger.setLevel(logging.INFO) + + +def create_file(path, contents): + """Writes contents to path if file doesn't exist, else raise + AssertionError + + """ + + if os.path.exists(path): + raise AssertionError('File {} already exists'.format(path)) + + with open(path, 'w') as outfile: + outfile.write(contents) + + +def write_file(path, contents): + """Writes contents to path if file""" + + with open(path, 'w') as outfile: + outfile.write(contents) + + +def read_json_file(path): + """Writes contents to path if file doesn't exist, else raise + AssertionError + + """ + + if not os.path.exists(path): + raise AssertionError('File {} does not exist'.format(path)) + + with open(path, 'r') as infile: + return json.loads(infile.read()) + + +def read_dir_json_files(path): + """Reads all files as json from directory""" + + for (dirpath, _, filenames) in os.walk(path): + for filename in filenames: + yield read_json_file(os.path.join(dirpath, filename)) + + return # only walk the first level + + +def merge_file_doc(existing_file_doc, file_doc): + """Merges the file doc into the existing file doc""" + + existing_case_subdocs = { + case_doc['case_id']: case_doc + for case_doc in existing_file_doc['cases'] + } + + this_case_subdocs = { + case_doc['case_id']: case_doc + for case_doc in file_doc['cases'] + } + + all_case_subdocs = dict(existing_case_subdocs, **this_case_subdocs) + existing_file_doc['cases'] = all_case_subdocs.values() + + return existing_file_doc + + +class GraphIndex(object): + """Base class to represent a complete or in progress index""" + + __metaclass__ = abc.ABCMeta + + @abc.abstractmethod + def __iter__(self): + """Iterate over the index (for backwards compatible unpacking)""" + + raise NotImplementedError + + @abc.abstractmethod + def case_doc_count(self): + """Returns the number of existing case docs""" + + raise NotImplementedError + + @abc.abstractmethod + def add_case_doc(self, case_doc): + """Adds a case document to this index""" + + raise NotImplementedError + + @abc.abstractmethod + def add_file_doc(self, file_doc): + """Adds a file document to this index""" + + raise NotImplementedError + + @abc.abstractmethod + def add_annotation_doc(self, annotation_doc): + """Adds a annotation document to this index""" + + raise NotImplementedError + + @abc.abstractmethod + def add_project_doc(self, case_doc): + """Adds a case document to this index""" + + raise NotImplementedError + + +class MemoryGraphIndex(GraphIndex): + """Class to represent a complete or in progress index on memory""" + + def __init__(self): + """Index constructor""" + + logger.info('Creating new %s', self) + + self._case_docs = {} + self._file_docs = {} + self._annotation_docs = {} + self._project_docs = {} + + def __iter__(self): + """Iterate over the index (for backwards compatible unpacking)""" + + return iter(( + self._case_docs.itervalues(), + self._file_docs.itervalues(), + self._annotation_docs.itervalues(), + self._project_docs.itervalues(), + )) + + def add_case_doc(self, case_doc): + """Adds a case document to this index""" + + case_id = case_doc['case_id'] + self._case_docs[case_id] = case_doc + + def add_file_doc(self, file_doc): + """Adds a file document to this index""" + + file_id = file_doc['file_id'] + + if file_id not in self._file_docs: + self._file_docs[file_id] = file_doc + + else: + self._file_docs[file_id] = merge_file_doc( + self._file_docs[file_id], file_doc) + + def add_annotation_doc(self, annotation_doc): + """Adds a annotation document to this index""" + + annotation_id = annotation_doc['annotation_id'] + self._annotation_docs.setdefault(annotation_id, annotation_doc) + + def add_project_doc(self, project_doc): + """Adds a project document to this index""" + + project_id = project_doc['project_id'] + self._project_docs[project_id] = project_doc + + def case_doc_count(self): + """Returns the number of existing case docs""" + + return len(self._case_docs) + + +class DiskGraphIndex(GraphIndex): + """Class to represent a complete or in progress index on memory""" + + + def __init__(self, data_dir_base=None): + """Index constructor""" + + rel_data_dir = '{}_{}'.format(data_dir_base, int(time.time())) + self.data_dir = os.path.abspath(os.path.expanduser(rel_data_dir)) + self.case_dir = os.path.join(self.data_dir, 'cases') + self.file_dir = os.path.join(self.data_dir, 'files') + self.annotation_dir = os.path.join(self.data_dir, 'annotations') + self.project_dir = os.path.join(self.data_dir, 'projects') + + self._seen_case_ids = set() + self._seen_file_ids = set() + + logger.info('Creating new %s', self) + + self.create_data_dir() + + def __repr__(self): + return "<{}('{}')>".format(self.__class__.__name__, self.data_dir) + + + def __iter__(self): + """Iterate over the index (for backwards compatible unpacking)""" + + return iter(( + read_dir_json_files(self.case_dir), + read_dir_json_files(self.file_dir), + read_dir_json_files(self.annotation_dir), + read_dir_json_files(self.project_dir), + )) + + def delete(self): + """Deletes the entire data_dir""" + + logger.info('Deleting %s', self) + shutil.rmtree(self.data_dir) + + def create_data_dir(self): + """Create directory to build index to disk""" + + if not os.path.exists(self.data_dir): + os.makedirs(self.data_dir) + + if not os.path.exists(self.case_dir): + os.makedirs(self.case_dir) + + if not os.path.exists(self.file_dir): + os.makedirs(self.file_dir) + + if not os.path.exists(self.annotation_dir): + os.makedirs(self.annotation_dir) + + if not os.path.exists(self.project_dir): + os.makedirs(self.project_dir) + + def add_case_doc(self, case_doc): + """Adds a case document to this index""" + + case_id = case_doc['case_id'] + path = os.path.join(self.case_dir, case_id) + self._seen_case_ids.add(case_id) + + create_file(path, json.dumps(case_doc)) + + def add_file_doc(self, file_doc): + """Adds a file document to this index""" + + file_id = file_doc['file_id'] + path = os.path.join(self.file_dir, file_id) + + try: + if file_id not in self._seen_file_ids: + return create_file(path, json.dumps(file_doc)) + + existing_file_doc = read_json_file(path) + updated_file_doc = merge_file_doc(existing_file_doc, file_doc) + + write_file(path, json.dumps(updated_file_doc)) + + except Exception as e: + print str(e) + logger.exception(e) + import pdb; pdb.set_trace() + raise + + finally: + self._seen_file_ids.add(file_id) + + + def add_annotation_doc(self, annotation_doc): + """Adds a annotation document to this index""" + + annotation_id = annotation_doc['annotation_id'] + path = os.path.join(self.annotation_dir, annotation_id) + + if not os.path.exists(path): + create_file(path, json.dumps(annotation_doc)) + + def add_project_doc(self, project_doc): + """Adds a project document to this index""" + + project_id = project_doc['project_id'] + path = os.path.join(self.project_dir, project_id) + + create_file(path, json.dumps(project_doc)) + + def case_doc_count(self): + """Returns the number of existing case docs""" + + return len(self._seen_case_ids) diff --git a/esbuild/graph/common/util.py b/esbuild/graph/common/util.py new file mode 100644 index 00000000..d05553d9 --- /dev/null +++ b/esbuild/graph/common/util.py @@ -0,0 +1,152 @@ +# -*- coding: utf-8 -*- +""" +esbuild.graph.common.util +---------------------------------- + +Common utilities for things like logging + +""" + +from datadog import statsd + +from progressbar import ( + ProgressBar, + Percentage, + Bar, + ETA, +) + + +def upsert_file_into_dict(files, file_doc): + """Merge this file document into all other relevant file documents (or + just add it if none exist) + + TODO: make this more descriptive + + """ + + did = file_doc['file_id'] + + if did not in files: + files[did] = file_doc + return + + for case in file_doc['cases']: + case_id = case['case_id'] + + existing_ids = { + case['case_id'] + for case in files[did]['cases'] + } + + if case_id not in existing_ids: + files[did]['cases'] += file_doc['cases'] + + +def get_file_to_case_paths(cls, file_to_case_paths): + """Given a node, return the paths the lead monotonically up to case""" + + parent_labels = { + link['dst_type'].label + for link in cls._pg_links.values() + } + return ( + path + for path in file_to_case_paths + if path and path[0] in parent_labels + ) + +def reverse_paths(paths, starting_label): + """Reverse the paths and add starting label to end of reversed paths""" + + return [ + list(reversed(l))[1:]+[starting_label] + for l in paths + ] + + +def remove_index_files(files, index_file_extensions): + """Returns a set of files that are not index files""" + + return { + file_ for file_ in files + if not is_index_file(file_, index_file_extensions) + } + + +def is_index_file(node, index_file_extensions): + """Given a node, return whether it is considerend an 'index file' + + :returns: bool + + """ + + # Active index files + if node._dictionary['category'] == 'index_file': + return True + + # Legacy index files + elif node.label == 'file': + for extension in index_file_extensions: + if node._props.get('file_name', '').endswith(extension): + return True + + else: + return False + + +def log_warning(logger, title, text, tags=[], *args, **kwargs): + """Log a warning to logger and statsd""" + + logger.warning("{}: {}".format(title, text)) + statsd.event( + title, + text, + source_type_name="esbuild", + alert_type="warning", + tags=tags, + ) + + +def log_error(logger, title, text, tags=[], *args, **kwargs): + """Log an error to logger and statsd""" + + logger.error("{}: {}".format(title, text)) + statsd.event( + title, + text, + source_type_name="esbuild", + alert_type="error", + tags=tags, + ) + + +def truncate_path(path, label): + """ + Given a path (a list of node labels), "truncate" it from the left + such that it starts with the given label, or return [], e.g.: + + truncate_path(["a", "b", "c"], "a") -> ["b", "c"] + truncate_path(["c", "d"], "b") -> [] + + """ + for i, currlabel in enumerate(path): + if currlabel == label: + return path[i+1:] + return [] + + +def get_pbar(title, maxval): + """Create and initialize a custom progressbar + + :param str title: The text of the progress bar + :param int maxval: The maximumum value of the progress bar + + """ + maxval = maxval or 1 # prevent maxal of 0 + pbar = ProgressBar(widgets=[ + title, Percentage(), ' ', + Bar(marker='#', left='[', right=']'), ' ', + ETA(), ' '], maxval=maxval) + pbar.update(0) + return pbar diff --git a/esbuild/graph/legacy/builder.py b/esbuild/graph/legacy/builder.py index 0fcdd6ad..61320087 100644 --- a/esbuild/graph/legacy/builder.py +++ b/esbuild/graph/legacy/builder.py @@ -86,7 +86,7 @@ def denormalize_archive_files(self, visited_file_ids=None): """ - archives = self.nodes_labeled('archive') + archives = self.cache.nodes_labeled('archive') file_docs = [] if visited_file_ids is None: @@ -95,7 +95,8 @@ def denormalize_archive_files(self, visited_file_ids=None): for archive in archives: file_docs.append(self.get_archive_as_file_doc(archive)) - for file_ in self.neighbors_labeled(archive, self.file_labels): + files = self.cache.neighbors_labeled(archive, self.file_labels) + for file_ in files: # skip any files visited in or before this function if file_.node_id in visited_file_ids: continue diff --git a/tests/conftest.py b/tests/conftest.py index b956c4a8..73cdcdaa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -39,6 +39,15 @@ _graph = PsqlGraphDriver(PG_HOST, PG_USER, PG_PASSWORD, PG_DATABASE) +@pytest.fixture(scope='session') +def psqlgraph_args(): + return ( + PG_HOST, + PG_USER, + PG_PASSWORD, + PG_DATABASE, + ) + @pytest.fixture def clear_database(): diff --git a/tests/test_graph_active_builder.py b/tests/test_graph_active_builder.py index c44afd7a..e75452be 100644 --- a/tests/test_graph_active_builder.py +++ b/tests/test_graph_active_builder.py @@ -25,6 +25,19 @@ subtree_paths_to_file, ) +from esbuild.graph.common.cache import ( + CachedGraph, +) + +from esbuild.graph.common.index import ( + DiskGraphIndex, +) + +from esbuild.graph.common.builder import ( + build_index, +) + + # Define the number of files that should be loaded as documents N_FILES = 9 @@ -33,23 +46,32 @@ @pytest.fixture(scope='module') -def index(): - builder = ActiveGraphIndexBuilder(_graph) - builder.cache_database() - index = builder.denormalize_all() - return Index._make(index) - - -@pytest.fixture -def cached_builder(scope='module'): - builder = ActiveGraphIndexBuilder(_graph) - builder.cache_database() - return builder - - -@pytest.fixture() -def builder(): - return ActiveGraphIndexBuilder(_graph) +def cached_graph(psqlgraph_args): + options = ActiveGraphIndexBuilder.get_caching_options() + cache = CachedGraph( + caching_options=options, + psqlgraph_driver_args=psqlgraph_args, + ) + cache.cache_database() + return cache + + +@pytest.yield_fixture(scope='module') +def index(psqlgraph_args): + disk_index = build_index( + ActiveGraphIndexBuilder, + psqlgraph_args, + '~/indexes', + ) + yield Index._make(map(list, disk_index)) + disk_index.delete() + + +@pytest.yield_fixture(scope='module') +def builder(cached_graph): + disk_index = DiskGraphIndex('~/indexes') + yield ActiveGraphIndexBuilder(cached_graph, disk_index) + disk_index.delete() @pytest.fixture @@ -333,6 +355,12 @@ def test_file_to_read_group_paths(label, path): assert path in ActiveGraphIndexBuilder.file_to_read_group_paths[label] +def lookup_expunged_node(builder, node): + return next( + expunged for expunged in builder.cache.graph.nodes() + if expunged.node_id == node.node_id + ) + def test_get_file_read_groups(graph, index): f_ids = {n.node_id for n in graph.nodes(md.SubmittedAlignedReads).all()} assert not [d for d in index.files if d['file_id'] in f_ids] @@ -342,24 +370,25 @@ def test_get_file_read_groups(graph, index): (md.AlignmentWorkflow, 2), (md.SomaticMutationCallingWorkflow, 2), ]) -def test_get_analysis_read_groups(graph, cached_builder, cls, count): +def test_get_analysis_read_groups(graph, builder, cls, count): for workflow in graph.nodes(cls).all(): - read_groups = list(cached_builder.get_analysis_read_groups(workflow)) + workflow = lookup_expunged_node(builder, workflow) + read_groups = list(builder.get_analysis_read_groups(workflow)) assert len(read_groups) == count for read_group in read_groups: assert read_group.label == 'read_group' - @pytest.mark.parametrize('cls,count', [ (md.AlignedReads, 1), (md.CopyNumberSegment, 1), (md.RunMetadata, 1), (md.ExperimentMetadata, 1), ]) -def test_get_file_associated_entities(graph, cached_builder, cls, count): +def test_get_file_associated_entities(graph, builder, cls, count): for node in graph.nodes(cls).all(): - if cached_builder.is_file_indexed(node): - entities = list(cached_builder.get_file_associated_entities(node)) + if builder.cache.is_file_indexed(node): + node = lookup_expunged_node(builder, node) + entities = list(builder.get_file_associated_entities(node)) assert len(entities) == count @@ -367,11 +396,12 @@ def test_get_file_associated_entities(graph, cached_builder, cls, count): (md.BiospecimenSupplement, 0), (md.ClinicalSupplement, 0), ], scope='module') -def test_add_related_files(graph, cached_builder, cls, count): +def test_add_related_files(graph, builder, cls, count): for node in graph.nodes(cls).all(): - if cached_builder.is_file_indexed(node): + if builder.cache.is_file_indexed(node): + node = lookup_expunged_node(builder, node) doc = {} - cached_builder.add_related_files(node, doc) + builder.add_related_files(node, doc) assert len(doc.get('metadata_files', [])) == count @@ -381,11 +411,12 @@ def test_add_related_files(graph, cached_builder, cls, count): (md.AlignedReads, False), (md.CopyNumberSegment, False), ], scope='module') -def test_add_archive(graph, cached_builder, cls, has_archive): +def test_add_archive(graph, builder, cls, has_archive): for node in graph.nodes(cls).all(): - if cached_builder.is_file_indexed(node): + if builder.cache.is_file_indexed(node): + node = lookup_expunged_node(builder, node) doc = {} - cached_builder.add_archives(node, doc) + builder.add_archives(node, doc) assert ('archive' in doc) == has_archive diff --git a/tests/test_graph_legacy_builder.py b/tests/test_graph_legacy_builder.py index 8c27025e..ab7f5169 100644 --- a/tests/test_graph_legacy_builder.py +++ b/tests/test_graph_legacy_builder.py @@ -10,6 +10,7 @@ from conftest import Index, _graph from data import fuzzed from esbuild.graph.legacy.builder import LegacyGraphIndexBuilder +from esbuild.graph.common.cache import CachedGraph from gdcdatamodel import models as md from jsonpath_rw import parse @@ -20,11 +21,16 @@ ) +def new_cache(graph=_graph): + options = LegacyGraphIndexBuilder.get_caching_options() + cache = CachedGraph(graph, options) + cache.cache_database() + return cache + + def build_index(graph): - builder = LegacyGraphIndexBuilder(graph) - builder.cache_database() - index = builder.denormalize_all() - return Index._make(index) + builder = LegacyGraphIndexBuilder(new_cache(graph)) + return Index._make(builder.denormalize_all()) # ====================================================================== @@ -33,7 +39,8 @@ def build_index(graph): @pytest.fixture() def builder(): - return LegacyGraphIndexBuilder(_graph) + return LegacyGraphIndexBuilder(new_cache()) + @pytest.fixture(scope="module") @@ -103,10 +110,14 @@ def test_path_value_in(index, doc_type, path, expected, count): assert actual.value in expected -def test_omitted_projects(graph): - builder = LegacyGraphIndexBuilder(graph) - builder.omitted_projects.add(('TCGA', 'BRCA')) - builder.cache_database() +def test_omitted_projects(graph, monkeypatch): + caching_options = LegacyGraphIndexBuilder.get_caching_options() + monkeypatch.setattr(caching_options, 'omitted_projects', + caching_options.omitted_projects.union({('TCGA', 'BRCA')})) + + cache = CachedGraph(graph, caching_options) + cache.cache_database() + builder = LegacyGraphIndexBuilder(cache) index = Index._make(builder.denormalize_all()) assert index.cases == [] @@ -146,7 +157,9 @@ def test_non_case_suppression(graph): redacted2.aliquots = [aliquot] s.add(redacted1) s.add(redacted2) + index = build_index(graph) + case_doc = [c for c in index.cases if c["case_id"] == case.node_id][0] sample_doc = [s for s in case_doc["samples"] if s["sample_id"] == sample.node_id][0] assert portion.node_id not in [ diff --git a/tests/test_suite_data.gv b/tests/test_suite_data.gv index 83c26ed5..e3adf706 100644 --- a/tests/test_suite_data.gv +++ b/tests/test_suite_data.gv @@ -2,7 +2,6 @@ digraph { graph [rankdir=RL] "5fa9998b-deff-493e-8a8e-dc2422192a48" [label=] "5fa9998b-deff-493e-8a8e-dc2422192a48" -> "eda6d2d5-4199-4f76-a45b-1d0401b4e54c" - "40407260-e805-4c2e-b2a7-13862bc5e494" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "008ba655-a0a3-42c4-8c72-f1341365ef02" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "aliquot-without-downstream" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "0ffb3f3d-f20e-43d1-9867-7dc75ac24f3b" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" @@ -11,6 +10,7 @@ digraph { "05c45162-6c94-4a15-accc-b6239451064c" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "7b017050-97d4-45bb-bf83-c89dab812e44" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "6d066a72-f59f-45a8-ab90-216000b36da4" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" + "40407260-e805-4c2e-b2a7-13862bc5e494" -> "5fa9998b-deff-493e-8a8e-dc2422192a48" "b3601406-3676-4f76-9aa0-ed68ed6c3a05" [label=] "b3601406-3676-4f76-9aa0-ed68ed6c3a05" -> "64f66bc3-1cee-41d7-ae86-cb443e84f30e" "973bd442-04a0-4189-8f02-c8c7e041afe9" -> "b3601406-3676-4f76-9aa0-ed68ed6c3a05" @@ -22,14 +22,14 @@ digraph { "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" [label=] "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" -> "c1e5beaa-6103-409d-bdd4-a86c0f210014" "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" -> "c8611490-4cbd-5651-8de2-64484a515eec" + "64f66bc3-1cee-41d7-ae86-cb443e84f30e" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" + "cnv-file-1" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "live-file" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "legacy-file-with-empty-acl" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "harmonized-file" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "non-live-file" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "to-delete-file" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "related-file" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" - "cnv-file-1" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" - "64f66bc3-1cee-41d7-ae86-cb443e84f30e" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "d7cb38ff-0ca2-5496-896b-92c5a76b6109" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" "rescinded-annotation" -> "84df0f82-69c4-4cd3-a4bd-f40d2d6ef916" biospecimen_supplement_1 [label=] @@ -50,10 +50,10 @@ digraph { "live-file" -> data_subtype_aligned_reads "live-file" -> "index-file" "live-file" -> "related-file" + "run-metadata-1" -> "live-file" + "analysis-metadata-1" -> "live-file" "experiment-metadata-1" -> "live-file" "harmonized-file" -> "live-file" - "analysis-metadata-1" -> "live-file" - "run-metadata-1" -> "live-file" "a2b74dcc-052a-42ce-836e-c2fb549beea5" [label=] "old-biospecimen-supplement-xml" [label=] "old-biospecimen-supplement-xml" -> "eda6d2d5-4199-4f76-a45b-1d0401b4e54c" @@ -119,6 +119,7 @@ digraph { "c1fd82a9-f75f-4297-b2c2-ec91c40a57f4" -> "956ca84c-1124-53ff-824f-fa0c84425425" "cnv-file-1" [label=] "cnv-workflow-1" -> "cnv-file-1" + "non-live-file" [label=] "snv-file" [label=] "snv-file" -> "eda6d2d5-4199-4f76-a45b-1d0401b4e54c" "bd4d1c78-c448-4bbf-8348-a77f3786c648" [label=] @@ -131,15 +132,15 @@ digraph { "64f66bc3-1cee-41d7-ae86-cb443e84f30e" [label=] "read-group-qc-1" -> "64f66bc3-1cee-41d7-ae86-cb443e84f30e" "40407260-e805-4c2e-b2a7-13862bc5e494" [label=] + "3013e9be-aa3e-4986-990c-559982f00e36" -> "40407260-e805-4c2e-b2a7-13862bc5e494" "3febc6c8-85ae-4d38-ba55-c959959846db" -> "40407260-e805-4c2e-b2a7-13862bc5e494" "5f5b9bb2-3278-424f-9cf2-e26f0c3b0fd5" -> "40407260-e805-4c2e-b2a7-13862bc5e494" - "3013e9be-aa3e-4986-990c-559982f00e36" -> "40407260-e805-4c2e-b2a7-13862bc5e494" "5f5b9bb2-3278-424f-9cf2-e26f0c3b0fd5" [label=] "6d066a72-f59f-45a8-ab90-216000b36da4" -> "5f5b9bb2-3278-424f-9cf2-e26f0c3b0fd5" "c1e5beaa-6103-409d-bdd4-a86c0f210014" [label=] "c1e5beaa-6103-409d-bdd4-a86c0f210014" -> "eda6d2d5-4199-4f76-a45b-1d0401b4e54c" - "5b2a99b7-e1a8-4739-acaf-d5f75cc47021" -> "c1e5beaa-6103-409d-bdd4-a86c0f210014" "0395a62f-3f37-4068-bab6-4c1d29cef2d5" -> "c1e5beaa-6103-409d-bdd4-a86c0f210014" + "5b2a99b7-e1a8-4739-acaf-d5f75cc47021" -> "c1e5beaa-6103-409d-bdd4-a86c0f210014" clinical_supplement_1 [label=] clinical_supplement_1 -> "eda6d2d5-4199-4f76-a45b-1d0401b4e54c" clinical_supplement_1 -> archive_1 @@ -157,7 +158,7 @@ digraph { "5b2a99b7-e1a8-4739-acaf-d5f75cc47021" [label=] "0395a62f-3f37-4068-bab6-4c1d29cef2d5" [label=] "analysis-metadata-1" [label=] - "non-live-file" [label=] + "index-file-2" [label=] "aggregated-somatic-mutation-1" [label=] "experiment-metadata-1" [label=] archive_1 [label=] @@ -189,7 +190,6 @@ digraph { "3febc6c8-85ae-4d38-ba55-c959959846db" [label=] "7b017050-97d4-45bb-bf83-c89dab812e44" -> "3febc6c8-85ae-4d38-ba55-c959959846db" "008ba655-a0a3-42c4-8c72-f1341365ef02" -> "3febc6c8-85ae-4d38-ba55-c959959846db" - "index-file-2" [label=] "008ba655-a0a3-42c4-8c72-f1341365ef02" [label=] "7ef3885b-37ce-5e16-8ba3-9d75b6690008" [label=] "7b017050-97d4-45bb-bf83-c89dab812e44" -> "7ef3885b-37ce-5e16-8ba3-9d75b6690008" diff --git a/tests/test_suite_data.gv.pdf b/tests/test_suite_data.gv.pdf index c467b687..fda71932 100644 Binary files a/tests/test_suite_data.gv.pdf and b/tests/test_suite_data.gv.pdf differ