diff --git a/csst_common/data_manager.py b/csst_common/data_manager.py index 1077852b9ec6b4119f3ead29dc554741f84422c1..e9dea92619004bb1c678fe237a56e28622499965 100644 --- a/csst_common/data_manager.py +++ b/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