add dm.cleanup and dm.n_jobs_gpu

csst_common/data_manager.py

@@ -111,7 +111,9 @@ class CsstMsDataManager:
     verbose : bool
         If True, print verbose info.
     n_jobs : int
-        The number of jobs.
+        The number of CPU jobs.
+    n_jobs_gpu : int
+        The number of GPU jobs.
     backend : str
         The joblib backend.
     device : str
@@ -162,6 +164,7 @@ class CsstMsDataManager:
                  clear_dir: bool = False,
                  verbose: bool = True,
                  n_jobs: int = 18,
+                 n_jobs_gpu: int = 1,
                  backend: str = "multiprocessing",
                  device: str = "CPU",
                  stamps: str = "",
@@ -228,7 +231,9 @@ class CsstMsDataManager:
         # record hard code names in history
         self.hardcode_history = []
 
+        # parallel configuration
         self.n_jobs = n_jobs
+        self.n_jobs_gpu = n_jobs_gpu
         self.backend = backend
         self.device = device
 
@@ -241,16 +246,30 @@ class CsstMsDataManager:
         if clear_dir:
             self.clear_dir(self.dir_l1)
 
+        # L1 whitelist
+        self.l1_whitelist = []
+        for detector in self.target_detectors:
+            if self.datatype == "mbi":
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="img_L1.fits"))
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="wht_L1.fits"))
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="flg_L1.fits"))
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="cat.fits"))
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="psf.fits"))
+            elif self.datatype == "sls":
+                self.l1_whitelist.append(self.l1_detector(detector=detector, post="L1_1.fits"))
+
         # pipeline logger
         if log_ppl == "":
             self.logger_ppl = get_logger(name="CSST L1 Pipeline Logger", filename="")
         else:
             self.logger_ppl = get_logger(name="CSST L1 Pipeline Logger", filename=os.path.join(dir_l1, log_ppl))
+            self.l1_whitelist.append(log_ppl)
         # module logger
         if log_mod == "":
             self.logger_mod = get_logger(name="CSST L1 Module Logger", filename="")
         else:
             self.logger_mod = get_logger(name="CSST L1 Module Logger", filename=os.path.join(dir_l1, log_mod))
+            self.l1_whitelist.append(log_mod)
 
         self.custom_bias = None
         self.custom_dark = None
@@ -891,6 +910,15 @@ class CsstMsDataManager:
         """ Query L1 data from DFS. """
         pass
 
+    def l1_cleanup(self):
+        filelist = glob.glob(f"{self.dir_l1}/**")
+        for file in filelist:
+            if file not in self.l1_whitelist:
+                try:
+                    os.remove(file)
+                except:
+                    pass
+
 
 # temporarily compatible with old interface
 CsstMbiDataManager = CsstMsDataManager