update data loading options

PurdueAF · Mar 4, 2024 · 4ecb6d1 · 4ecb6d1
1 parent 4cdfa68
commit 4ecb6d1
Show file tree

Hide file tree

Showing 15 changed files with 42 additions and 42 deletions.
diff --git a/af_benchmark/data_access/loader.py b/af_benchmark/data_access/loader.py
@@ -4,31 +4,39 @@
 
 def get_file_list(cls):
     mode = cls.config.get('data-access.mode', 'local')
-    if mode == 'local':
+    if mode == 'explicit-files':
         file_list = cls.config.get('data-access.files', [])
-    elif mode == 'local_dirs':
+    elif mode == 'explicit-dirs':
         dirs = cls.config.get('data-access.directories', [])
         file_list = []
         for dir in dirs:
             file_list.extend(glob.glob(dir+"/**/*.root", recursive = True))
-    elif mode == 'dbs_dataset':
-        dbsdataset = cls.config.get('data-access.dataset', "")
+    elif mode == 'dbs-dataset':
+        dbsdatasets = cls.config.get('data-access.datasets', [])
         xrootdserver = cls.config.get('data-access.xrootdserver', 'eos.cms.rcac.purdue.edu:1094')
         dbs = DbsApi('https://cmsweb.cern.ch/dbs/prod/global/DBSReader')
-        file_list = ["root://"+xrootdserver+"/"+file['logical_file_name'] for file in dbs.listFiles(dataset=dbsdataset)]
-    elif mode == 'dbs_block':
-        dbsblock = cls.config.get('data-access.block', "")
+        file_list = [
+            "root://"+xrootdserver+"/"+file['logical_file_name']
+                for dataset in dbsdatasets
+                for file in dbs.listFiles(dataset=dataset)
+        ]
+    elif mode == 'dbs-block':
+        dbsblocks = cls.config.get('data-access.blocks', [])
         xrootdserver = cls.config.get('data-access.xrootdserver', 'eos.cms.rcac.purdue.edu:1094')
         dbs = DbsApi('https://cmsweb.cern.ch/dbs/prod/global/DBSReader')
-        file_list = ["root://"+xrootdserver+"/"+file['logical_file_name'] for file in dbs.listFiles(block_name=dbsblock)]
-    elif mode == 'dbs_file':
-        dbsfile = cls.config.get('data-access.file', "")
+        file_list = [
+            "root://"+xrootdserver+"/"+file['logical_file_name']
+                for block in dbsblocks
+                for file in dbs.listFiles(block_name=block)
+        ]
+    elif mode == 'dbs-files':
+        dbsfiles = cls.config.get('data-access.files', [])
         xrootdserver = cls.config.get('data-access.xrootdserver', 'cms-xcache.rcac.purdue.edu:1094')
         dbs = DbsApi('https://cmsweb.cern.ch/dbs/prod/global/DBSReader')
-        file_list = ["root://"+xrootdserver+"/"+dbsfile]
+        file_list = ["root://"+xrootdserver+"/"+file for file in dbsfiles]
     else:
         raise NotImplementedError(
-            f"Data access modes other than 'local' and 'local_dir' are not yet implemented"
+            f"Data access mode {mode} not implemented"
         )
 
     cls.n_files = len(file_list)

diff --git a/af_benchmark/example-configs/example-config-1.yaml b/af_benchmark/example-configs/example-config-1.yaml
diff --git a/...ple-configs/example-config-dbs-block.yaml → ...le-configs/example-config-dbs-blocks.yaml b/...ple-configs/example-config-dbs-block.yaml → ...le-configs/example-config-dbs-blocks.yaml
@@ -1,11 +1,12 @@
 executor:
   backend: sequential
 data-access:
-  mode: dbs_block
+  mode: dbs-block
   xrootdserver: eos.cms.rcac.purdue.edu:1094
 #  xrootdserver: cms-xcache.rcac.purdue.edu:1094
 #  xrootdserver: cmsxrootd.hep.wisc.edu:1094
-  block: /SingleMuon/Run2017C-02Apr2020-v1/NANOAOD#44236284-ff8d-4b95-b971-dcec15b5130f        
+  blocks:
+      - /SingleMuon/Run2017C-02Apr2020-v1/NANOAOD#44236284-ff8d-4b95-b971-dcec15b5130f        
 processor:
   parallelize_over: files
   columns: 5

diff --git a/...e-configs/example-config-dbs-dataset.yaml → ...-configs/example-config-dbs-datasets.yaml b/...e-configs/example-config-dbs-dataset.yaml → ...-configs/example-config-dbs-datasets.yaml
@@ -1,11 +1,12 @@
 executor:
-  backend: sequential
+  backend: futures
 data-access:
-  mode: dbs_dataset
+  mode: dbs-dataset
 #  xrootdserver: eos.cms.rcac.purdue.edu:1094
   xrootdserver: cms-xcache.rcac.purdue.edu:1094
-  dataset: /SingleMuon/Run2016B-02Apr2020_ver2-v1/NANOAOD
+  datasets:
+      - /SingleMuon/Run2016B-02Apr2020_ver2-v1/NANOAOD
 processor:
   parallelize_over: files
-  columns: 5
-  operation: sum
+  columns: 1
+  operation: nothing
diff --git a/...mple-configs/example-config-dbs-file.yaml → ...ple-configs/example-config-dbs-files.yaml b/...mple-configs/example-config-dbs-file.yaml → ...ple-configs/example-config-dbs-files.yaml
@@ -1,11 +1,12 @@
 executor:
   backend: sequential
 data-access:
-  mode: dbs_file
+  mode: dbs-files
 #  xrootdserver: eos.cms.rcac.purdue.edu:1094
   xrootdserver: cms-xcache.rcac.purdue.edu:1094
 #  xrootdserver: cmsxrootd.hep.wisc.edu:1094
-  file: /store/data/Run2016B/SingleMuon/NANOAOD/02Apr2020_ver2-v1/20000/90322FC2-4027-0E47-92E4-22307EC8EAD2.root       
+  files:
+      - /store/data/Run2016B/SingleMuon/NANOAOD/02Apr2020_ver2-v1/20000/90322FC2-4027-0E47-92E4-22307EC8EAD2.root       
 processor:
   parallelize_over: columns
 #  parallelize_over: files

diff --git a/...le-configs/example-config-local-dirs.yaml → ...configs/example-config-explicit-dirs.yaml b/...le-configs/example-config-local-dirs.yaml → ...configs/example-config-explicit-dirs.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: sequential
 data-access:
-  mode: local_dirs
+  mode: explicit-dirs
   directories:
       - /eos/purdue/store/data/Run2016B/SingleMuon/NANOAOD/02Apr2020_ver2-v1/
 processor:

diff --git a/af_benchmark/example-config.yaml → ...hmark/example-configs/example-config.yaml b/af_benchmark/example-config.yaml → ...hmark/example-configs/example-config.yaml
@@ -4,7 +4,7 @@ executor:
   # backend: dask-local
   # backend: dask-gateway
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
   # mode: local_dirs

diff --git a/af_benchmark/processor/uproot_processor.py b/af_benchmark/processor/uproot_processor.py
@@ -74,7 +74,7 @@ def process_column(self, file, column, **kwargs):
     def run_operation(self, column_data, **kwargs):
         operation = self.config.get('processor.operation', None)
 
-        if not operation:
+        if (not operation) or (operation=='nothing'):
             return
 
         data_in_memory = np.array([])

diff --git a/tests/test-dask-local.yaml b/tests/test-dask-local.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: dask-local
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-default.yaml b/tests/test-default.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: futures
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-explicit-columns.yaml b/tests/test-explicit-columns.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: sequential
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-futures.yaml b/tests/test-futures.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: futures
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-read-by-column.yaml b/tests/test-read-by-column.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: futures
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-read-by-file.yaml b/tests/test-read-by-file.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: futures
 data-access:
-  mode: local
+  mode: explicit-files
   files: 
     - tests/data/nano_dimuon.root
 processor:

diff --git a/tests/test-read-from-dir.yaml b/tests/test-read-from-dir.yaml
@@ -1,7 +1,7 @@
 executor:
   backend: futures
 data-access:
-  mode: local_dirs
+  mode: explicit-dirs
   directories: 
     - tests/data/
 processor: