Fix liftover caching bug in tests (#878)

* env vars

* remove liftover from hail

* ruff

* spelling

* move remove_liftover

* set value on mock

* ruff

* move from before tasks to after tasks

* move it back

* Update base_hail_table.py

* Update base_hail_table.py

* ws

Author: bpblanken

.github/workflows/unit-tests.yml

@@ -36,6 +36,8 @@ jobs:
         run: ruff . --output-format github
       - name: Unit Tests
         run: |
+          export GRCH37_TO_GRCH38_LIFTOVER_REF_PATH=v03_pipeline/var/test/liftover/grch37_to_grch38.over.chain.gz
+          export GRCH38_TO_GRCH37_LIFTOVER_REF_PATH=v03_pipeline/var/test/liftover/grch38_to_grch37.over.chain.gz
           export ACCESS_PRIVATE_REFERENCE_DATASETS=1
           export PYSPARK_SUBMIT_ARGS='--driver-memory 8G pyspark-shell'
           nosetests --with-coverage --cover-package v03_pipeline/lib v03_pipeline/lib

v03_pipeline/lib/annotations/liftover.py

@@ -15,3 +15,12 @@ def add_rg37_liftover(grch37_to_grch38_liftover_ref_path: str) -> None:
     rg38 = hl.get_reference(ReferenceGenome.GRCh38.value)
     if not rg37.has_liftover(rg38):
         rg37.add_liftover(grch37_to_grch38_liftover_ref_path, rg38)
+
+
+def remove_liftover():
+    rg37 = hl.get_reference(ReferenceGenome.GRCh37.value)
+    rg38 = hl.get_reference(ReferenceGenome.GRCh38.value)
+    if rg37.has_liftover(rg38):
+        rg37.remove_liftover(rg38)
+    if rg38.has_liftover(rg37):
+        rg38.remove_liftover(rg37)

v03_pipeline/lib/tasks/base/base_hail_table.py

@@ -1,6 +1,7 @@
 import hail as hl
 import luigi
 
+from v03_pipeline.lib.annotations.liftover import remove_liftover
 from v03_pipeline.lib.logger import get_logger
 from v03_pipeline.lib.model import Env
 from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
@@ -27,6 +28,13 @@ def init_hail(self):
         # Interval ref data join causes shuffle death, this prevents it
         hl._set_flags(use_new_shuffle='1', no_whole_stage_codegen='1')  # noqa: SLF001
 
+        # Ensure any cached liftover files within Hail are cleared
+        # to provide a clean context free of hidden state.
+        # This runs "before" a task to account for situations where
+        # the Hail write fails and we do not have the chance to
+        # run this method in the "after".
+        remove_liftover()
+
 
 # NB: these are defined over luigi.Task instead of the BaseHailTableTask so that
 # they work on file dependencies.

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples_test.py

@@ -44,6 +44,9 @@
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 from v03_pipeline.var.test.vep.mock_vep_data import MOCK_37_VEP_DATA, MOCK_38_VEP_DATA
 
+GRCH38_TO_GRCH37_LIFTOVER_REF_PATH = (
+    'v03_pipeline/var/test/liftover/grch38_to_grch37.over.chain.gz'
+)
 TEST_MITO_MT = 'v03_pipeline/var/test/callsets/mito_1.mt'
 TEST_SNV_INDEL_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
 TEST_SV_VCF = 'v03_pipeline/var/test/callsets/sv_1.vcf'
@@ -248,6 +251,7 @@ def test_multiple_update_vat(
             {'ENST00000327044': 'NM_015658.4'},
         )
         # make register_alleles return CAIDs for 4 of 30 variants
+        mock_env.GRCH38_TO_GRCH37_LIFTOVER_REF_PATH = GRCH38_TO_GRCH37_LIFTOVER_REF_PATH
         mock_env.SHOULD_REGISTER_ALLELES = True
         mock_register_alleles.side_effect = [
             iter(

v03_pipeline/lib/tasks/write_new_variants_table.py

@@ -12,7 +12,10 @@
 from v03_pipeline.lib.misc.callsets import get_callset_ht
 from v03_pipeline.lib.misc.io import remap_pedigree_hash
 from v03_pipeline.lib.misc.math import constrain
-from v03_pipeline.lib.model import Env, ReferenceDatasetCollection
+from v03_pipeline.lib.model import (
+    Env,
+    ReferenceDatasetCollection,
+)
 from v03_pipeline.lib.paths import (
     new_variants_table_path,
     variant_annotations_table_path,
@@ -214,7 +217,6 @@ def create_table(self) -> hl.Table:
             ):
                 ar_ht = ar_ht.union(ar_ht_chunk)
             new_variants_ht = new_variants_ht.join(ar_ht, 'left')
-
         return new_variants_ht.select_globals(
             updates={
                 hl.Struct(