hack at hdfs

v03_pipeline/lib/reference_datasets/clinvar.py

@@ -12,7 +12,7 @@
     CLINVAR_PATHOGENICITIES_LOOKUP,
 )
 from v03_pipeline.lib.model.definitions import ReferenceGenome
-from v03_pipeline.lib.reference_datasets.misc import vcf_to_ht
+from v03_pipeline.lib.reference_datasets.misc import safely_add_to_hdfs, vcf_to_ht
 
 CLINVAR_GOLD_STARS_LOOKUP = hl.dict(
     {
@@ -109,6 +109,7 @@ def get_submission_summary_ht() -> hl.Table:
         timeout=10,
     ) as r:
         shutil.copyfileobj(r.raw, tmp_file)
+        safely_add_to_hdfs(tmp_file.name)
     ht = hl.import_table(
         tmp_file.name,
         force=True,
@@ -163,6 +164,7 @@ def get_ht(
         delete=False,
     ) as tmp_file, requests.get(clinvar_url, stream=True, timeout=10) as r:
         shutil.copyfileobj(r.raw, tmp_file)
+        safely_add_to_hdfs(tmp_file.name)
     ht = vcf_to_ht(tmp_file.name, reference_genome)
     submitters_ht = get_submission_summary_ht()
     ht = ht.annotate(

v03_pipeline/lib/reference_datasets/dbnsfp.py

@@ -59,7 +59,7 @@ def get_ht(path: str, reference_genome: ReferenceGenome) -> hl.Table:
     types = TYPES[reference_genome]
     rename = RENAME[reference_genome]
 
-    with download_zip_file(path) as unzipped_dir:
+    with download_zip_file(path, 'dbnsfp') as unzipped_dir:
         ht = hl.import_table(
             f'{unzipped_dir}/dbNSFP*_variant.chr*.gz',
             types=types,

v03_pipeline/lib/reference_datasets/helix_mito.py

@@ -5,6 +5,7 @@
 import requests
 
 from v03_pipeline.lib.model.definitions import ReferenceGenome
+from v03_pipeline.lib.reference_datasets.misc import safely_add_to_hdfs
 
 RENAME = {
     'counts_hom': 'AC_hom',
@@ -22,6 +23,7 @@ def get_ht(
         delete=False,
     ) as tmp_file, requests.get(url, stream=True, timeout=10) as r:
         shutil.copyfileobj(r.raw, tmp_file)
+        safely_add_to_hdfs(tmp_file.name)
     ht = hl.import_table(
         tmp_file.name,
         types={

v03_pipeline/lib/reference_datasets/local_constraint_mito.py

@@ -9,7 +9,9 @@
 
 
 def get_ht(url: str, reference_genome: ReferenceGenome) -> hl.Table:
-    with download_zip_file(url, decode_content=True) as unzipped_dir:
+    with download_zip_file(
+        url, 'local_constraint_mito', decode_content=True
+    ) as unzipped_dir:
         ht = hl.import_table(
             os.path.join(
                 unzipped_dir,

v03_pipeline/lib/reference_datasets/misc.py

@@ -1,5 +1,6 @@
 import contextlib
 import os
+import subprocess
 import tempfile
 import zipfile
 
@@ -118,14 +119,16 @@ def copyfileobj(fsrc, fdst, decode_content, length=16 * 1024):
 
 
 @contextlib.contextmanager
-def download_zip_file(url, suffix='.zip', decode_content=False):
+def download_zip_file(url, dataset_name: str, suffix='.zip', decode_content=False):
+    dir_ = f'/tmp/{dataset_name}'  # noqa: S108
     with tempfile.NamedTemporaryFile(
+        dir=dir_,
         suffix=suffix,
     ) as tmp_file, requests.get(url, stream=True, timeout=10) as r:
         copyfileobj(r.raw, tmp_file, decode_content)
         with zipfile.ZipFile(tmp_file.name, 'r') as zipf:
             zipf.extractall(os.path.dirname(tmp_file.name))
-        # Extracting the zip file
+        safely_add_to_hdfs(dir_)
         yield os.path.dirname(tmp_file.name)
 
 
@@ -148,3 +151,21 @@ def select_for_interval_reference_dataset(
         **additional_selects,
     )
     return ht.key_by('interval')
+
+
+def safely_add_to_hdfs(file_name: str):
+    if os.getenv('HAIL_DATAPROC') != '1':
+        return
+    subprocess.run(
+        [  # noqa: S603
+            '/usr/bin/hdfs',
+            'dfs',
+            '-copyFromLocal',
+            '-f',
+            f'file://{file_name}',
+            file_name,
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )

v03_pipeline/lib/reference_datasets/mitimpact.py

@@ -11,7 +11,7 @@ def get_ht(
     reference_genome: ReferenceGenome,
 ) -> hl.Table:
     extracted_filename = url.removesuffix('.zip').split('/')[-1]
-    with download_zip_file(url, suffix='.txt.zip') as unzipped_dir:
+    with download_zip_file(url, 'mitimpact', suffix='.txt.zip') as unzipped_dir:
         ht = hl.import_table(
             os.path.join(
                 unzipped_dir,

v03_pipeline/lib/reference_datasets/screen.py

@@ -6,6 +6,7 @@
 
 from v03_pipeline.lib.model import ReferenceGenome
 from v03_pipeline.lib.reference_datasets.misc import (
+    safely_add_to_hdfs,
     select_for_interval_reference_dataset,
 )
 
@@ -20,6 +21,7 @@ def get_ht(path: str, reference_genome: ReferenceGenome) -> hl.Table:
         timeout=10,
     ) as r:
         shutil.copyfileobj(r.raw, tmp_file)
+        safely_add_to_hdfs(tmp_file.name)
     ht = hl.import_table(
         tmp_file.name,
         no_header=True,