From d68a00f46e6c273af9b10a20965b23e246a3f97a Mon Sep 17 00:00:00 2001
From: lruizcalico <lruiz@calicolabs.com>
Date: Fri, 27 Oct 2023 13:31:43 -0700
Subject: [PATCH] add helper h5

---
 src/borzoi/helpers/h5_grad_utils.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/borzoi/helpers/h5_grad_utils.py b/src/borzoi/helpers/h5_grad_utils.py
index ed55af7..4e1f96f 100644
--- a/src/borzoi/helpers/h5_grad_utils.py
+++ b/src/borzoi/helpers/h5_grad_utils.py
@@ -1,6 +1,7 @@
 import h5py
 import numpy as np
 import argparse
+import subprocess
 import tempfile
 import os
 from baskerville.helpers.gcs_utils import download_from_gcs, upload_file_gcs
@@ -67,9 +68,8 @@ def collect_h5_borzoi(out_dir, num_procs, sad_stat) -> None:
 
 def download_h5_gcs(output_gcs_dir, num_processes) -> str:
     temp_dir = tempfile.mkdtemp()  # create a temp dir for output
-    out_dir = temp_dir + "/output_dir"
-    print(f"outdir is {out_dir}")
-    print(f"gcs_out_dir is {output_gcs_dir}")
+    print(f"temp_dir is {temp_dir}")
+    out_dir = temp_dir + "/" + output_gcs_dir.split("/")[-1]
     if not os.path.isdir(out_dir):
         os.mkdir(out_dir)
     # download output from tempfile
@@ -81,6 +81,12 @@ def download_h5_gcs(output_gcs_dir, num_processes) -> str:
             f"{out_dir}/job{pi}/scores_f0c0.h5",
         )
         print(f"Done downloading {pi} partition")
+    # download all of the files in the folder
+    # Use gsutil to copy the contents recursively
+    # subprocess.check_call(["gsutil", "-m", "cp", "-r", output_gcs_dir, temp_dir])
+    # out_dir = temp_dir + "/" + output_gcs_dir.split("/")[-1]
+    print(f"outdir is {out_dir}")
+    print(f"gcs_out_dir is {output_gcs_dir}")
     print(f"Done dowloading")
     return out_dir
 
@@ -103,7 +109,7 @@ def main():
     if args.gcs:
         # download files to tempdir
         local_out_dir = download_h5_gcs(args.out_dir, args.num_procs)
-    collect_h5_borzoi(args.file_name, local_out_dir, args.num_procs)
+    collect_h5_borzoi(local_out_dir, args.num_procs, args.sad_stat)
     # upload to gcs
     print(f"is there such a file? {local_out_dir}/scores_f0c0.h5")
     print(os.path.isfile(f"{local_out_dir}/scores_f0c0.h5"))