hail-is / hail

Cloud-native genomic dataframes and batch computing
https://hail.is
MIT License
968 stars 242 forks source link

[query] Hana / SEQR need support optimizing Hail Query code #13882

Open danking opened 10 months ago

danking commented 10 months ago

What happened?

Hana Snow is the engineer for SEQR.

Previously, SEQR used elastic search as its datastore. Unfortunately, elastic search was very expensive because, to get reasonable performance, SEQR indexed nearly every field. The ES index was huge and the VM resources necessary to run an ES instance on that index were expensive (like 1000s USD per month).

I've been supporting Hana as much as I can, but she needs someone who can be more dedicated and responsive than me.

She uses a k8s cluster. She has a SEQR frontend deployment. She also has a Hail deployment (statefulset maybe?). The Hail pod has an SSD mounted read-only. That SSD has all the SEQR data in Hail Table form. There are many tables with annotations (variant metadata, like "probability this variant is damaging" or "likely causes this to happen to the protein"). There are also "per-family" tables which contain all the sequences within a single family. Many queries are directly against a particular family. Those tables are small and quick to read.

There's also one giant table containing all the sequences from all the families. That table is large and expensive to read. A lot of our engineering work has been around making sure queries against that table are fast.

Tim, at one point, had enough of her system locally that he could experiment with running queries on his laptop against his SSD. He hacked on the queries themselves and on Hail itself until the bandwidth was fast enough that the queries should complete fast enough on the full dataset. Fast enough varies but generally a couple tens of seconds is OK.

The work here is to pair with Hana to diagnose performance issues and make changes until the queries are acceptably fast. The first thing I would do is update her to the latest Hail (with the array decoder improvement as well as the memory overhead stuff on which Daniel is working). Then, with Hana's help, test the timing of some queries. If the queries are still too slow, your options are:

  1. Check the log files and the IR. Are there unnecessary shuffles? Is the code really large? Can we do less work maybe?
  2. Have Hana help you replicate her setup locally. You just need a slice of the data and enough of SEQR to run a query. Now hook up a profiler. What's slow? Can we do something about that?

Version

0.2.124

Relevant log output

No response

danking commented 10 months ago

Dan to connect Hana with Ed. Reassign once that's done.

hanars commented 10 months ago

One obvious point of optimization that Dan already identified is the way we were keying our data is causing full shuffles, as we were keying the data by a string variant ID in the form 1-32683987-ACTCTT-A instead of locus-allele. Changing back to keying on locus-allele fixes this issue for our more straightforward searches, but we have a search that looks for pairs of possible compound heterozygous variants in the same gene, and that still is resulting in 2 full shuffles. I'm a little at a loss for how to fix this, because we are grouping by an unsorted field so I'm not sure how to prevent us from working with an unsorted dataset. The offending code right now is as follows (somewhat simplified for readability):

primary_variants = hl.agg.filter(ch_ht[HAS_ALLOWED_ANNOTATION], hl.agg.collect(ch_ht.row))
secondary_variants = hl.agg.filter(ch_ht[HAS_ALLOWED_SECONDARY_ANNOTATION], hl.agg.collect(ch_ht.row))
ch_ht = ch_ht.group_by('gene_ids').aggregate(v1=primary_variants, v2=secondary_variants)
ch_ht = ch_ht.explode(ch_ht.v1)
ch_ht = ch_ht.explode(ch_ht.v2)
ch_ht = ch_ht.annotate(grouped_variants=hl.sorted([ch_ht.v1, ch_ht.v2], key=lambda v: (v.locus, v.alleles)))
ch_ht = ch_ht.key_by(
  locus=ch_ht.grouped_variants[0].locus, 
  alleles=ch_ht.grouped_variants[0].alleles,
  locus2=ch_ht.grouped_variants[1].locus, 
  alleles2=ch_ht.grouped_variants[1].alleles,
)
ch_ht = ch_ht.distinct()
...
# more filtering and annotating
...
return ch_ht._key_by_assert_sorted('locus', 'alleles')
danking commented 10 months ago

Hmm. OK, this is an interesting one. You're basically trying to do:

  1. For each gene-interval, get all the variants in that gene.
  2. Filter to variants having the primary or secondary allowed annotation.
  3. Compute the crossproduct of variants matching the primary annotation and variants matching the secondary annotation.
  4. Filter the cross-product to the upper-right triangle including the diagonal (aka only store one of (x, y) and (y, x)).
  5. More filtering.

Let me see what the team thinks.

hanars commented 10 months ago

Filter the cross-product to the upper-right triangle including the diagonal (aka only store one of (x, y) and (y, x)).

To be a little pedantic, this filtering is not always filtering to the upper right triangle, as the primary and secondary variants are not necessarily the same set, and are sometimes even fully distinct. But "only store one of (x, y) and (y, x)" is accurate

hanars commented 10 months ago

Just to update, when updating to 0.2.125 I discovered a bug :( https://hail.zulipchat.com/#narrow/stream/123010-Hail-Query-0.2E2-support/topic/array.20index.20out.20of.20bounds.20error.20on.20dict.2Eget

hanars commented 10 months ago

we now have our dev test environment running with hail 0.2.126 and this query took ~92 seconds. So faster, but we do still need some performance enhancements. @ehigham let me know what would be helpful for you to get you started on this effort

danking commented 10 months ago

@hanars , for my own curiosity, what was the timing on this query before 0.2.126?

hanars commented 10 months ago

104s. So like a 10% improvement

hanars commented 10 months ago

I was able to figure out how to remove all the "network shuffle"s and "coerced sorted dataset"s and that improved the time down to 73 seconds, so a big improvement! I would still hope to improve performance a bit more, being reliably under a minute would be helpful

Here are the logs from that search, let me know what else I can do to help improve the performance or to help you figure it out: hail-search.log

PR is here if you are interested: https://github.com/broadinstitute/seqr/pull/3717

danking commented 10 months ago

cc: @ehigham

hanars commented 10 months ago

seqr hail search code is here: https://github.com/broadinstitute/seqr/tree/master/hail_search

Running the aiohttp service is just python -m hail_search. Runs on hail 0.2.126

The data you need is gs://seqr-datasets/v03/GRCh38/SNV_INDEL/runs/manual__2023-11-07T23-31-23.149902+00-00 For running the service, you need a DATASETS_DIR env variable defined, and that data should be available at $DATASETS_DIR/GRCh38/SNV_INDEL

Post body for a relatively quick search:

{
  "genome_version": "GRCh38",
  "num_results": 100,
  "annotations": {
    "in_frame": [
      "inframe_insertion",
      "inframe_deletion"
    ],
    "missense": [
      "stop_lost",
      "initiator_codon_variant",
      "start_lost",
      "protein_altering_variant",
      "missense_variant"
    ],
    "nonsense": [
      "stop_gained"
    ],
    "splice_ai": "0.2",
    "frameshift": [
      "frameshift_variant"
    ],
    "structural": [],
    "extended_splice_site": [],
    "essential_splice_site": [
      "splice_donor_variant",
      "splice_acceptor_variant"
    ],
    "structural_consequence": [
      "LOF",
      "DUP_LOF",
      "INV_SPAN",
      "COPY_GAIN"
    ]
  },
  "datasetType": "VARIANTS",
  "pathogenicity": {
    "hgmd": [
      "disease_causing"
    ],
    "clinvar": [
      "pathogenic",
      "likely_pathogenic"
    ]
  },
  "dataset_type": "ALL",
  "secondary_dataset_type": null,
  "inheritance_mode": "de_novo",
  "inheritance_filter": {
    "A": "has_alt",
    "N": "ref_ref"
  },
  "sample_data": {
    "SNV_INDEL": [
      {
        "sample_id": "RGP_2436_2_D1",
        "individual_guid": "I0097169_rgp_2436_2",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "N"
      },
      {
        "sample_id": "RGP_2436_3_D1",
        "individual_guid": "I0097170_rgp_2436_3",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "A"
      },
      {
        "sample_id": "RGP_2436_1_D1",
        "individual_guid": "I0097168_rgp_2436_1",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "N"
      }
    ]
  },
  "sort": "xpos",
  "sort_metadata": null,
  "frequencies": {
    "g1k": {
      "ac": null,
      "af": 1
    },
    "exac": {
      "ac": null,
      "af": 1
    },
    "topmed": {
      "ac": null,
      "af": 1
    },
    "gnomad_svs": {
      "ac": null,
      "af": 0.001
    },
    "sv_callset": {
      "ac": null,
      "af": 0.001
    },
    "gnomad_exomes": {
      "ac": null,
      "af": 0.001
    },
    "gnomad_genomes": {
      "ac": null,
      "af": 0.001
    },
    "seqr": {
      "ac": null,
      "af": 0.01
    }
  },
  "quality_filter": {
    "min_ab": 20,
    "min_gq": 30,
    "min_qs": 50,
    "vcf_filter": "pass"
  },
  "custom_query": null,
  "intervals": null,
  "exclude_intervals": null,
  "gene_ids": null,
  "variant_ids": null,
  "rs_ids": null
}

Example of a slow search that needs improvement:

{
  "genome_version": "GRCh38",
  "num_results": 100,
  "annotations": {
    "other": [
      "non_coding_exon_variant"
    ],
    "in_frame": [
      "inframe_insertion",
      "inframe_deletion"
    ],
    "missense": [
      "stop_lost",
      "initiator_codon_variant",
      "start_lost",
      "protein_altering_variant",
      "missense_variant"
    ],
    "nonsense": [
      "stop_gained"
    ],
    "splice_ai": "0.1",
    "frameshift": [
      "frameshift_variant"
    ],
    "structural": [
      "gCNV_DUP",
      "gCNV_DEL"
    ],
    "synonymous": [],
    "extended_splice_site": [],
    "essential_splice_site": [
      "splice_donor_variant",
      "splice_acceptor_variant"
    ],
    "structural_consequence": [
      "LOF",
      "DUP_LOF",
      "MSV_EXON_OVR",
      "INTRAGENIC_EXON_DUP",
      "PARTIAL_EXON_DUP"
    ]
  },
  "pathogenicity": {
    "hgmd": [
      "disease_causing"
    ],
    "clinvar": [
      "pathogenic",
      "likely_pathogenic",
      "vus_or_conflicting"
    ]
  },
  "annotations_secondary": {
    "other": [
      "transcript_ablation",
      "transcript_amplification",
      "5_prime_UTR_variant",
      "3_prime_UTR_variant",
      "TFBS_ablation",
      "TFBS_amplification",
      "TF_binding_site_variant",
      "regulatory_region_variant",
      "regulatory_region_ablation",
      "regulatory_region_amplification",
      "non_coding_transcript_exon_variant__canonical"
    ],
    "in_frame": [
      "inframe_insertion",
      "inframe_deletion"
    ],
    "missense": [
      "stop_lost",
      "initiator_codon_variant",
      "start_lost",
      "protein_altering_variant",
      "missense_variant"
    ],
    "nonsense": [
      "stop_gained"
    ],
    "frameshift": [
      "frameshift_variant"
    ],
    "structural": [
      "gCNV_DEL",
      "gCNV_DUP"
    ],
    "synonymous": [
      "synonymous_variant",
      "stop_retained_variant"
    ],
    "extended_splice_site": [
      "splice_region_variant"
    ],
    "essential_splice_site": [
      "splice_donor_variant",
      "splice_acceptor_variant"
    ],
    "structural_consequence": [
      "LOF",
      "DUP_LOF",
      "INTRONIC",
      "UTR",
      "PROMOTER",
      "INTRAGENIC_EXON_DUP",
      "PARTIAL_EXON_DUP"
    ]
  },
  "dataset_type": "ALL",
  "secondary_dataset_type": "ALL",
  "inheritance_mode": "recessive",
  "inheritance_filter": {
    "A": null,
    "N": null
  },
  "sample_data": {
    "SNV_INDEL": [
      {
        "sample_id": "RGP_2436_2_D1",
        "individual_guid": "I0097169_rgp_2436_2",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "N"
      },
      {
        "sample_id": "RGP_2436_3_D1",
        "individual_guid": "I0097170_rgp_2436_3",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "A"
      },
      {
        "sample_id": "RGP_2436_1_D1",
        "individual_guid": "I0097168_rgp_2436_1",
        "family_guid": "F041731_rgp_2436",
        "project_guid": "R0594_rare_genomes_project_gen",
        "affected": "N"
      }
    ]
  },
  "sort": "xpos",
  "sort_metadata": null,
  "frequencies": {
    "g1k": {
      "ac": null,
      "af": 1
    },
    "exac": {
      "ac": null,
      "af": 1
    },
    "topmed": {
      "ac": null,
      "af": 1
    },
    "gnomad_svs": {
      "ac": null,
      "af": 0.01
    },
    "sv_callset": {
      "ac": null,
      "af": 0.01
    },
    "gnomad_exomes": {
      "ac": null,
      "af": 0.01
    },
    "gnomad_genomes": {
      "ac": null,
      "af": 0.01
    },
    "seqr": {
      "ac": null,
      "af": 0.03
    }
  },
  "quality_filter": {
    "min_ab": 10,
    "min_gq": 40,
    "min_qs": 50
  },
  "custom_query": null,
  "intervals": null,
  "exclude_intervals": null,
  "gene_ids": null,
  "variant_ids": null,
  "rs_ids": null
}

Let me know if there is anything else you need that I can help with!

ehigham commented 9 months ago

Here are the mean elapsed times for the two queries above as measured on my machine using the local and spark backends with localised data.

query local spark
0 7s 5s
1 20s 13s

@hanars, do these relative differences between queries match your observations?

hanars commented 9 months ago

In our dev environment the first query runs in 10s, but the second one runs in 77s, if it ran in 20 we'd be fine with it. As a sanity check, the first query should have 3 results and the second should have 85

danking commented 9 months ago

One possible point of confusion, there's a Hail Query "LocalBackend" which is a mostly internal thing and that's different than spark-in-local-mode (aka hl.init(..., master='local[*]')); I'm pretty sure we're using spark-in-local-mode for SEQR things.

ehigham commented 9 months ago

Ah the timings were slightly off as I had not downloaded all the data and was using some from hail_search/fixtures. After pulling down the SNV_INDELS data, my updated timings are:

query results elapsed
0 4 7s
1 83 50s
hanars commented 9 months ago

I'm pretty sure we're using spark-in-local-mode for SEQR things.

We run hl.init(idempotent=True) on the hail docker image for seqr. Not sure whether that is running spark-in-local mode or not

hanars commented 9 months ago

the data should be really deterministic so I'm not sure why those counts are off, but at least those counts and times are more accurate to what we are seeing so probably it means your test environment is reasonable. I would recommend against using the hail_search/fixtures data for anything other than running the test suite, its not designed to play nicely with real data

ehigham commented 9 months ago

the data should be really deterministic so I'm not sure why those counts are off, but at least those counts and times are more accurate to what we are seeing so probably it means your test environment is reasonable.

I guess the exact query isn't important, more the code shape and the kinds of operations hail's doing. Given the timings, it may be more-or-less representative.

ehigham commented 9 months ago

Dan and I have profiled the second query and have a couple of thoughts.

We've identified one source of slowness in the compiler and generated code that relates to how we handle I/O. This is currently under active development and hope will lead to decent performance improvements. The change itself is significant, however, so I can't comment on timelines for when to expect the work to be complete by.

There may be some code tweaks that do less work as discussed. I'll have a play with the source code in hail_search and report back.

hanars commented 9 months ago

awesome, thanks for looking into it!

hanars commented 9 months ago

The current logic of key_by and distinct to remove duplicate pairs is actually not properly deduplicating pairs in some cases. Since I know that thats not really the approach we want in the long run I didn't bother figuring out why, and instead tried implementing a version of the code that filters out duplicate pairs before exploding, which should result in less work overall. However, I somehow introduced an extra Ordering unsorted dataset with network shuffle and the overall performance of that search got slower by 15 seconds.

Here is the change I made. Let me know if you have a better approach for filtering out duplicates, or if you see any ways to reorganize this code to make it less shuffle-y https://github.com/broadinstitute/seqr/commit/2e45403efc159b58cec723f86e6de7653d64cf5f

ehigham commented 9 months ago

I've taking a similar approach to try to remove unwanted elements before exploding. Sadly, I haven't seen any noticable improvement. I'm also not sure about correctness as I couldn't get your changes in 2e45403 to work (got errors about missing gene_ids). I saw ~20 fewer results in the second query with the code below.

Here's what I wrote based off master.

    def _filter_compound_hets(self):
        ch_ht = self._ht
        if self._is_recessive_search:
            ch_ht = ch_ht.filter(ch_ht.comp_het_family_entries.any(hl.is_defined))

        # Get possible pairs of variants within the same gene
        ch_ht = ch_ht.annotate(gene_ids=self._gene_ids_expr(ch_ht, comp_het=True))
        ch_ht = ch_ht.explode(ch_ht.gene_ids)

        # Filter allowed transcripts to the grouped gene
        transcript_annotations = {
            k: ch_ht[k].filter(lambda t: t.gene_id == ch_ht.gene_ids)
            for k in [ALLOWED_TRANSCRIPTS, ALLOWED_SECONDARY_TRANSCRIPTS] if k in ch_ht.row
        }
        if transcript_annotations:
            ch_ht = ch_ht.annotate(**transcript_annotations)
        primary_filters = self._get_annotation_filters(ch_ht)
        secondary_filters = self._get_annotation_filters(ch_ht, is_secondary=True)

        self.unfiltered_comp_het_ht = ch_ht.filter(hl.any(primary_filters + secondary_filters))
        if self._has_secondary_annotations and not (primary_filters and secondary_filters):
            # In cases where comp het pairs must have different data types, there are no single data type results
            return None

        primary_variants = hl.agg.filter(hl.any(primary_filters), hl.agg.collect(ch_ht.row))
        if secondary_filters:
            row_agg = ch_ht.row
            if ALLOWED_TRANSCRIPTS in row_agg and ALLOWED_SECONDARY_TRANSCRIPTS in row_agg:
                # Ensure main transcripts are properly selected for primary/secondary annotations in variant pairs
                row_agg = row_agg.annotate(**{ALLOWED_TRANSCRIPTS: row_agg[ALLOWED_SECONDARY_TRANSCRIPTS]})
            secondary_variants = hl.agg.filter(hl.any(secondary_filters), hl.agg.collect(row_agg))
        else:
            secondary_variants = primary_variants

>>> start of changes:
        def unique_combinations(v1s, v2s):
            return hl.rbind(
                hl.sorted(
                    hl.array(hl.set(v2s).difference(hl.set(v1s))),
                    key=lambda x: x.variant_id
                ),
                lambda v2s_uniq:
                    v1s.flatmap(lambda v1:
                        v2s_uniq.map(lambda v2: hl.array([v1, v2]))
                    )
            )

        ch_ht = ch_ht \
            .group_by('gene_ids') \
            .aggregate(vs=unique_combinations(primary_variants, secondary_variants)) \
            .explode('vs') \
            .key_by()

        ch_ht = ch_ht.annotate(
            valid_families=hl.map(
                self._is_valid_comp_het_family,
                ch_ht.vs[0].comp_het_family_entries,
                ch_ht.vs[1].comp_het_family_entries
            )
        )
        ch_ht = ch_ht.filter(ch_ht.valid_families.any(lambda x: x))
        return ch_ht.select(**{
            GROUPED_VARIANTS_FIELD: hl.array([
                self._annotated_comp_het_variant(ch_ht, ch_ht.vs[k])
                for k in [0, 1]
            ])
        })

    @staticmethod
    def _annotated_comp_het_variant(ch_ht, variant):
        return variant.annotate(
            gene_id=ch_ht.gene_ids,
            family_entries=hl.enumerate(ch_ht.valid_families).filter(
                lambda x: x[1]).map(lambda x: variant.comp_het_family_entries[x[0]]),
        )
hanars commented 9 months ago

Oh yeah, to get my code to work you need to comment out line 778 gene_id=ch_ht.gene_ids, in _annotated_comp_het_variant. It doesn't break search to be missing that annotation, it just has some downstream display affects that I would need to fix if I actually wanted to use the code, but given the performance hit I wasn't sure it was worth figuring that out as this code may be too slow to use

I was not able to get the code you provided here to run either, but one concern I have with it is that the unique combinations are computed per gene, but if you have a pair of variants that are each in the same 2 genes, you would get the pair twice in the results, one for each gene

The error I get when I run the code you provide is

"Key type mismatch: cannot index table with given expressions:
  Table key:         <<<empty key>>>
  Index Expressions: locus<GRCh38>, array<str>, set<str>, array<array<struct{GQ: int32, AB: float64, DP: int32, GT: call, sampleId: str, sampleType: str, individualGuid: str, familyGuid: str, affected_id: int32}>>, array<array<struct{GQ: int32, AB: float64, DP: int32, GT: call, sampleId: str, sampleType: str, individualGuid: str, familyGuid: str, affected_id: int32}>>, struct{z_score: float32}, struct{region_type_ids: array<int32>}, locus<GRCh37>, str, array<struct{amino_acids: str, canonical: int32, codons: str, gene_id: str, hgvsc: str, hgvsp: str, transcript_id: str, biotype_id: int32, consequence_term_ids: array<int32>, is_lof_nagnag: bool, lof_filter_ids: array<int32>, transcript_rank: int32}>, str, int64, struct{PHRED: float32}, struct{alleleId: int32, conflictingPathogenicities: array<struct{pathogenicity_id: int32, count: int32}>, goldStars: int32, pathogenicity_id: int32, assertion_ids: array<int32>}, struct{REVEL_score: float32, VEST4_score: float32, MutPred_score: float32, SIFT_pred_id: int32, Polyphen2_HVAR_pred_id: int32, MutationTaster_pred_id: int32, fathmm_MKL_coding_pred_id: int32}, struct{Eigen_phred: float32}, struct{AF_POPMAX: float32, AF: float32, AC_Adj: int32, AC_Het: int32, AC_Hom: int32, AC_Hemi: int32, AN_Adj: int32}, struct{AF: float32, AN: int32, AC: int32, Hom: int32, AF_POPMAX_OR_GLOBAL: float32, FAF_AF: float32, Hemi: int32}, struct{AF: float32, AN: int32, AC: int32, Hom: int32, AF_POPMAX_OR_GLOBAL: float32, FAF_AF: float32, Hemi: int32}, struct{MPC: float32}, struct{score: float32}, struct{delta_score: float32, splice_consequence_id: int32}, struct{AC: int32, AF: float32, AN: int32, Hom: int32, Het: int32}, struct{accession: str, class_id: int32}, struct{AC: int32, AN: int32, AF: float32, hom: int32}, array<struct{amino_acids: str, canonical: int32, codons: str, gene_id: str, hgvsc: str, hgvsp: str, transcript_id: str, biotype_id: int32, consequence_term_ids: array<int32>, is_lof_nagnag: bool, lof_filter_ids: array<int32>, transcript_rank: int32}>, bool, array<struct{amino_acids: str, canonical: int32, codons: str, gene_id: str, hgvsc: str, hgvsp: str, transcript_id: str, biotype_id: int32, consequence_term_ids: array<int32>, is_lof_nagnag: bool, lof_filter_ids: array<int32>, transcript_rank: int32}>, bool, str"
ehigham commented 9 months ago

You'll have to provide me with the full error to help, I don't know where that's coming from. Perhaps I edited another function elsewhere and didn't include it up there.

that are each in the same 2 genes

The same two genes? How can a gene appear twice when you've grouped by the gene id?

hanars commented 9 months ago

So lets say primary variants are A and B and the secondary variants are B and C. And lets say we have genes 1 and 2. Lets say the table grouped by gene ids looks like:

gene | primary | secondary
1 | A, B | B, C
2 | B    | B, C

The possible unique permutations by gene would then be

gene | primary | secondary
1 | A | B
1 | A | C
1 | B | C
2 | B | C

So all the permutations by gene are unique, but then when you just pull out the pairs themselves the pair of variant B, C would be returned twice

ehigham commented 9 months ago

You're concerned about having the same pair across different genes? That's what the code in master did, right? Which gene do you want the pair to be associated with? Why not add a group-by pair to get them all?

hanars commented 9 months ago

yeah, it keyed by the pairs and then called distinct, which should theoretically get you the list of distinct pairs. Bit it did that at the end after computing and annotating all the possibilities

ehigham commented 9 months ago

Oh I see. Thanks for clarifying - I wasn't sure what that bit did! That should be a simple fix, though perhaps at this point not worth it as this is not a fruitful optimisation for this query

danking commented 9 months ago

Next steps:

  1. upload the profile, the mt.describe(), metadata.json.gz from the MT/HT to the team chat and get feedback (Chris, Patrick take a look).

Decode appears quite slow.

hanars commented 9 months ago

That should be a simple fix, though perhaps at this point not worth it as this is not a fruitful optimization for this query

Agreed, although depending on the time line for a good optimization for this query I may circle back on this, as there is currently a bug in this deduplication so if the actual optimizaion won;t be available for a while it might be worth fixing in the meantime

hanars commented 9 months ago

Hey, I know everything at the Broad pretty much grinds to a halt in December, but just wanted to check in on where things are before we head into the break!

ehigham commented 8 months ago

Sorry, I've been on my honeymoon and just got back. I resume work on this in the new year. Apologies for the delay.

ehigham commented 7 months ago

Ok, getting back up to speed on this. There've been a number of changes on either side of this project, so going to give new timings and profiling results for the two queries.

Here are mean timings for the two queries, run 5 times, and taking the mean of all but the first. It seems there's been a considerable regression since November on the second query, highlighting our need to get automated benchmark runs in per release (https://github.com/hail-is/hail/issues/14221).

query spark
0 7s
1 87s

Attached are YourKit profiler results of the two queries. 'fast' refers to query 0 and 'slow' to the longer-running query 1. seqr-profile-data.zip seqr-logs.zip

hanars commented 7 months ago

FWIW, I think most of the regression was this PR: https://github.com/broadinstitute/seqr/pull/3792/files

We kind of knew the performance was worse but also we decided that we needed to fix the bug more than we needed performance to be good :/

hanars commented 7 months ago

We recently put in a couple PRs that improve performance on these searches so thought I would update here. They were mostly changes to things upstream of the portion of code we have been focusing on and change how data is initially read in, but the biggest performance gain we got was adding hl._set_flags(use_new_shuffle='1'). A lot of the focus was around how we handle searches in multiple data types which has been out of the scope of this work so far, so for the search we've been profiling here its only came down to like 80 seconds, but figured its worth sharing. Hopefully this does not cause to catastrophic of a merge conflict for you guys

https://github.com/broadinstitute/seqr/pull/3873 https://github.com/broadinstitute/seqr/pull/3876

ehigham commented 7 months ago

Updated baseline timings on 61a5d4834 with Hana's updates + use_new_shuffle=1 on my pc:

--------------------------------------
Timing query 'fast' with 5 repeats.
Initial:   8.920063795001624s                                 
   Mean:   6.641188349500226s
--------------------------------------
Timing query 'slow' with 5 repeats.
Initial:   45.03917969699978s 
   Mean:   42.99382913374939s
--------------------------------------
hanars commented 7 months ago

Awesome, thats really exciting! Are you sure that the right commit? That one seems to be about dropping hail-az support...

ehigham commented 7 months ago

That's just the commit I'm currently using for later reference