From e66ed791bd9ac01a68739156e8555ca2d05316db Mon Sep 17 00:00:00 2001
From: BO ZHANG <bozhang@nao.cas.cn>
Date: Fri, 23 Dec 2022 09:48:09 +0800
Subject: [PATCH] add n_jobs and backend to dm

---
 csst_common/data_manager.py | 67 +++++++++++++++++++++++++------------
 1 file changed, 45 insertions(+), 22 deletions(-)

diff --git a/csst_common/data_manager.py b/csst_common/data_manager.py
index 6b29960..e55f01b 100644
--- a/csst_common/data_manager.py
+++ b/csst_common/data_manager.py
@@ -107,6 +107,10 @@ class CsstMsDataManager:
         If True, clear ``dm.dir_l1`` directory.
     verbose : bool
         If True, print verbose info.
+    n_jobs : int
+        The number of jobs.
+    backend : str
+        The joblib backend.
 
     Examples
     --------
@@ -152,6 +156,8 @@ class CsstMsDataManager:
                  log_mod="csst-l1mod.log",
                  clear_dir=False,
                  verbose=True,
+                 n_jobs=18,
+                 backend="multiprocessing"
                  ):
 
         # version
@@ -211,7 +217,8 @@ class CsstMsDataManager:
         # record hard code names in history
         self.hardcode_history = []
 
-        self.n_jobs = 18
+        self.n_jobs = n_jobs
+        self.backend = backend
 
         # aXe
         self.set_env()
@@ -287,7 +294,17 @@ class CsstMsDataManager:
         raise DeprecationWarning("This method is deprecated, please directly use dm.target_detectors = detectors!")
 
     @staticmethod
-    def from_dir(ver_sim="C5.2", datatype="mbi", dir_l0=".", dir_l1=".", path_aux="", use_dfs=False, dfs_node="kmust"):
+    def from_dir(
+            ver_sim="C5.2",
+            datatype="mbi",
+            dir_l0=".",
+            dir_l1=".",
+            path_aux="",
+            use_dfs=False,
+            dfs_node="kmust",
+            n_jobs=18,
+            backend="multiprocessing"
+    ):
         """ initialize the multi-band imaging data manager """
 
         assert ver_sim in ["C5.2", ]
@@ -310,24 +327,27 @@ class CsstMsDataManager:
         # exp_stop = int(exp_stop)
         # obs_id = int(obs_id)
 
-        return CsstMsDataManager(ver_sim=ver_sim,
-                                 datatype=datatype,
-                                 available_detectors=available_detectors,
-                                 target_detectors=None,
-                                 dir_l0=dir_l0,
-                                 dir_l1=dir_l1,
-                                 path_aux=path_aux,  # bias dark flat
-                                 use_dfs=use_dfs,
-                                 dfs_node=dfs_node,
-                                 obs_id=obs_id,
-                                 exp_start=exp_start,
-                                 exp_stop=exp_stop,
-                                 _telescope=_telescope,
-                                 _instrument=_instrument,
-                                 _survey=_survey,
-                                 obs_type=obs_type,
-                                 l0_post="_".join(l0_post),
-                                 )
+        return CsstMsDataManager(
+            ver_sim=ver_sim,
+            datatype=datatype,
+            available_detectors=available_detectors,
+            target_detectors=None,
+            dir_l0=dir_l0,
+            dir_l1=dir_l1,
+            path_aux=path_aux,  # bias dark flat
+            use_dfs=use_dfs,
+            dfs_node=dfs_node,
+            obs_id=obs_id,
+            exp_start=exp_start,
+            exp_stop=exp_stop,
+            _telescope=_telescope,
+            _instrument=_instrument,
+            _survey=_survey,
+            obs_type=obs_type,
+            l0_post="_".join(l0_post),
+            n_jobs=n_jobs,
+            backend=backend
+        )
 
     @staticmethod
     def glob_image(dir_l0, ver_sim="C5.2"):
@@ -641,7 +661,9 @@ class CsstMsDataManager:
             use_dfs=True,
             dfs_node="pml",
             clear_l1=False,
-            dfs_root="/share/dfs"
+            dfs_root="/share/dfs",
+            n_jobs=18,
+            backend="multiprocessing"
     ):
         """ Initialize CsstMsDataManager from DFS. """
         # (clear and) make directories
@@ -684,7 +706,8 @@ class CsstMsDataManager:
         # initialize dm
         dm = CsstMsDataManager.from_dir(
             ver_sim="C5.2", datatype=datatype, dir_l0=dir_l0, dir_l1=dir_l1,
-            path_aux=path_aux, use_dfs=use_dfs, dfs_node=dfs_node
+            path_aux=path_aux, use_dfs=use_dfs, dfs_node=dfs_node,
+            n_jobs=n_jobs, backend=backend
         )
         assert dm.obs_id == obs_id
 
-- 
GitLab