From 9b6cdb0346e43c87e430265afaea0ae282339503 Mon Sep 17 00:00:00 2001
From: kremin <kremin@umich.edu>
Date: Mon, 25 Nov 2024 13:46:11 -0800
Subject: [PATCH] remove temporary preproc directories in nightqa

---
 bin/desi_night_qa       |  3 ++-
 py/desispec/night_qa.py | 20 +++++++++++++++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/bin/desi_night_qa b/bin/desi_night_qa
index bfd72e651..194e39d3a 100755
--- a/bin/desi_night_qa
+++ b/bin/desi_night_qa
@@ -125,7 +125,8 @@ def main():
         log.info("args.compute_missing_only set: will override args.steps={}".format(args.steps))
         steps_tbd = []
         for step in steps_all:
-            if not os.path.isfile(outfns[step]) and step != 'html':
+            ## recompute html since recomputing other steps
+            if not os.path.isfile(outfns[step]) or step == 'html':
                 steps_tbd.append(step)
                 log.info(f"add {step} to steps_tbd, as args.compute_missing_only "
                          + f"set and no {outfns[step]} present")
diff --git a/py/desispec/night_qa.py b/py/desispec/night_qa.py
index ae9b3b8b8..e1239e8e1 100644
--- a/py/desispec/night_qa.py
+++ b/py/desispec/night_qa.py
@@ -8,6 +8,7 @@
 import os
 from glob import glob
 import tempfile
+import shutil
 import textwrap
 from desiutil.log import get_logger
 import multiprocessing
@@ -535,9 +536,11 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
                 run_preproc = False
         # AR run preproc?
 
+    temp_dir_loc = None
     if run_preproc:
         cmds = []
-        specprod_dir = tempfile.mkdtemp()
+        temp_dir_loc = tempfile.mkdtemp()
+        specprod_dir = temp_dir_loc
         outdir = os.path.dirname(findfile("preproc", night, dark_expid,
                                           'r1', specprod_dir=specprod_dir))
         os.makedirs(outdir, exist_ok=True)
@@ -550,7 +553,7 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
             # like os.system(cmd), but avoids system call for MPI compatibility
             cmds.append(cmd.split()[1:])
         with multiprocessing.Pool(processes=nproc) as pool:
-            pool.starmap(preproc.main, cmds)
+            pool.map(preproc.main, cmds)
     else:
         specprod_dir = prod
 
@@ -573,6 +576,11 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
     with multiprocessing.Pool(processes=nproc) as pool:
         mydicts = pool.starmap(_read_dark, myargs)
 
+    ## Remove the temporary directory
+    if temp_dir_loc is not None and os.path.exists(temp_dir_loc):
+            shutil.rmtree(temp_dir_loc)
+            log.info(f"Removed temporary directory and its contents: {temp_dir_loc}")
+            
     # AR campets to be additionally processed with --bkgsub-for-science
     # AR besides, check the very unlikely case where a camera is missing
     # AR    for all petals
@@ -594,6 +602,7 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
         bkgsub_science_cameras_str = ",".join(bkgsub_science_cameras)
         log.info("campets to be additionally processed with --bkgsub-for-science: {}".format(bkgsub_science_campets))
         #
+        bkgsub_specprod_dir = None
         if len(bkgsub_science_campets) > 0:
             bkgsub_specprod_dir = tempfile.mkdtemp()
             outdir = os.path.dirname(findfile("preproc", night, dark_expid,
@@ -609,7 +618,7 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
                 # like os.system(cmd), but avoids system call for MPI compatibility
                 cmds.append(cmd.split()[1:])
             with multiprocessing.Pool(processes=nproc) as pool:
-                pool.starmap(preproc.main, cmds)
+                pool.map(preproc.main, cmds)
 
         # AR read the bkgsub dark
         bkgsub_myargs = []
@@ -630,6 +639,11 @@ def create_dark_pdf(outpdf, night, prod, dark_expid, nproc, binning=4, bkgsub_sc
         with multiprocessing.Pool(processes=nproc) as pool:
             bkgsub_mydicts = pool.starmap(_read_dark, bkgsub_myargs)
 
+        ## Remove the temporary directory
+        if bkgsub_specprod_dir is not None and os.path.exists(bkgsub_specprod_dir):
+            shutil.rmtree(bkgsub_specprod_dir)
+            log.info(f"Removed temporary directory and its contents: {bkgsub_specprod_dir}")
+    
     # AR plotting
     # AR remarks:
     # AR - the (x,y) conversions for the side panels