checkpoint doesnt partition

Author: bpblanken

v03_pipeline/lib/misc/io.py

@@ -273,31 +273,26 @@ def remap_pedigree_hash(remap_path: str, pedigree_path: str) -> hl.Int32Expressi
 
 def checkpoint(
     t: hl.Table | hl.MatrixTable,
-    repartition_factor: int = 1,
 ) -> tuple[hl.Table | hl.MatrixTable, str]:
     suffix = 'mt' if isinstance(t, hl.MatrixTable) else 'ht'
     read_fn = hl.read_matrix_table if isinstance(t, hl.MatrixTable) else hl.read_table
     checkpoint_path = os.path.join(
         Env.HAIL_TMP_DIR,
         f'{uuid.uuid4()}.{suffix}',
     )
-    t.write(checkpoint_path, repartition_factor=repartition_factor)
+    t.write(checkpoint_path)
     return read_fn(checkpoint_path), checkpoint_path
 
 
 def write(
     t: hl.Table | hl.MatrixTable,
     destination_path: str,
     repartition: bool = True,
-    # May be used to increase the number of partitions beyond
-    # the optimally computed number.  A higher number will
-    # shard the table into more partitions.
-    repartition_factor: int = 1,
 ) -> hl.Table | hl.MatrixTable:
     t, path = checkpoint(t)
     if repartition:
         t = t.repartition(
-            (compute_hail_n_partitions(file_size_bytes(path)) * repartition_factor),
+            compute_hail_n_partitions(file_size_bytes(path)),
             shuffle=False,
         )
     return t.write(destination_path, overwrite=True)

v03_pipeline/lib/reference_datasets/splice_ai.py

@@ -1,6 +1,10 @@
 import hail as hl
 
-from v03_pipeline.lib.misc.io import checkpoint
+from v03_pipeline.lib.misc.io import (
+    checkpoint,
+    compute_hail_n_partitions,
+    file_size_bytes,
+)
 from v03_pipeline.lib.model import ReferenceGenome
 from v03_pipeline.lib.reference_datasets.misc import vcf_to_ht
 
@@ -13,7 +17,12 @@ def get_ht(
     # of file descriptors on dataproc :/
     hl._set_flags(use_new_shuffle=None, no_whole_stage_codegen='1')  # noqa: SLF001
     ht = vcf_to_ht(paths, reference_genome)
-    ht, _ = checkpoint(ht, repartition_factor=2)
+    ht, checkpoint_path = checkpoint(ht)
+    # The default partitions are too big, leading to OOMs.
+    ht = ht.repartition(
+        compute_hail_n_partitions(file_size_bytes(checkpoint_path)),
+        shuffle=False,
+    )
 
     # SpliceAI INFO field description from the VCF header: SpliceAIv1.3 variant annotation. These include
     # delta scores (DS) and delta positions (DP) for acceptor gain (AG), acceptor loss (AL), donor gain (DG), and