triggerflow 0.2__tar.gz → 0.2.2__tar.gz

{triggerflow-0.2 → triggerflow-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: triggerflow
-Version: 0.2
+Version: 0.2.2
 Summary: Utilities for ML models targeting hardware triggers
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License

{triggerflow-0.2 → triggerflow-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "triggerflow"
-version = "0.2"
+version = "0.2.2"
 description = "Utilities for ML models targeting hardware triggers"
 readme = "README.md"
 requires-python = ">=3.11"

{triggerflow-0.2 → triggerflow-0.2.2}/src/trigger_loader/cluster_manager.py

@@ -4,9 +4,6 @@ import logging
 from typing import Any
 
 from dask.distributed import Client, LocalCluster
-from dask_cuda import LocalCUDACluster
-from dask_jobqueue import HTCondorCluster
-from dask_kubernetes import KubeCluster
 
 logger = logging.getLogger(__name__)
 
@@ -63,15 +60,18 @@ class ClusterManager:
             self.cluster = LocalCluster(**self.cluster_config)
 
         elif ct == "condor":
+            from dask_jobqueue import HTCondorCluster
             self.cluster = HTCondorCluster(**self.cluster_config)
             if self.jobs and self.jobs > 0:
                 # Scale to the requested number of jobs
                 self.cluster.scale(jobs=self.jobs)
 
         elif ct == "cuda":
+            from dask_cuda import LocalCUDACluster
             self.cluster = LocalCUDACluster(**self.cluster_config)
 
         elif ct == "kubernetes":
+            from dask_kubernetes import KubeCluster
             self.cluster = KubeCluster(**self.cluster_config)
             if self.jobs and self.jobs > 0:
                 try:

{triggerflow-0.2 → triggerflow-0.2.2}/src/trigger_loader/loader.py

@@ -45,8 +45,58 @@ class TriggerLoader:
         )
 
     def _load_sample_json(self, sample_json: str) -> dict:
+        """
+        Loads the JSON and resolves file paths using the priority:
+        1. Explicit 'files' list or directory path (Local/Explicit)
+        2. 'DAS' query (Remote Fallback)
+
+        Returns the canonical coffea fileset format: {dataset_name: [file_path_list]}.
+        """
+        import glob
+        import os
+        
+        # Helper function definition needed here if it's not imported:
+        # def _fetch_files_from_das(das_query: str) -> list[str]: ... (placeholder or actual implementation)
+
         with open(sample_json) as f:
-            return json.load(f)
+            full_data = json.load(f)
+            dataset_metadata = full_data.get("samples", full_data)
+
+        fileset = {}
+        for ds_name, ds_info in dataset_metadata.items():
+            files = []
+            
+            if "files" in ds_info:
+                file_info = ds_info["files"]
+                
+                if isinstance(file_info, list):
+                    files = file_info
+                    
+                elif isinstance(file_info, str):
+                    if os.path.isdir(file_info):
+                        path_glob = os.path.join(file_info, "*.root")
+                        files = glob.glob(path_glob)
+                        logger.info(f"Resolved {len(files)} files from directory {file_info}.")
+                    else:
+                        files = [file_info]
+                
+                if files:
+                    logger.info(f"Using {len(files)} local/explicit files for {ds_name}.")
+ 
+            if not files and "DAS" in ds_info:
+                try:
+                    files = _fetch_files_from_das(ds_info["DAS"]) 
+                    logger.info(f"Resolved {len(files)} files via DAS for {ds_name}.")
+                except NameError:
+                    logger.error("DAS fetching skipped: _fetch_files_from_das is not defined.")
+
+            if not files:
+                logger.warning(f"No files found for dataset: {ds_name}. Skipping.")
+                continue
+
+            fileset[ds_name] = files
+
+        return fileset
 
     def _write_run_metadata_file(self, path: str, duration_s: float | None = None):
         meta_path = f"{path}/run_metadata.json"
@@ -58,9 +108,11 @@ class TriggerLoader:
             json.dump(data, f, indent=2)
 
     def _run(self, runner: processor.Runner, label: str):
-        logger.log(f"Starting processing ({label})...")
+        logger.log(logging.INFO, f"Starting processing ({label})...")
         start = time.time()
         proc = self._build_processor()
+        print(self.fileset)
+        
         acc = runner(
             self.fileset,
             treename="Events",
@@ -68,7 +120,7 @@ class TriggerLoader:
         )
         elapsed = time.time() - start
         self._write_run_metadata_file(self.output_path, elapsed)
-        logger.log(f"Finished in {elapsed:.2f}s (run_uuid={self.run_uuid})")
+        logger.log(logging.INFO, f"Finished in {elapsed:.2f}s (run_uuid={self.run_uuid})")
         return acc
 
     def run_distributed(self, cluster_type: str, cluster_config: dict,