triggerflow 0.1.11__py3-none-any.whl → 0.2__py3-none-any.whl

trigger_dataset/core.py

@@ -0,0 +1,88 @@
+import warnings
+from abc import ABC, abstractmethod
+from fnmatch import filter as fnmatch_filter
+
+import pandas as pd
+import uproot
+
+
+class TriggerDataset(ABC):
+    """
+    Abstract Base Class for loading data from ROOT files.
+
+    Users must inherit from this class and implement the abstract methods.
+    The core processing logic in `process_file` is fixed and cannot be overridden.
+    """
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def get_features(self) -> list[str]:
+        """
+        Return a list of branch names or patterns to keep from the dataset.
+        Accepts wildcards (e.g. Jet_*).
+        """
+        pass
+
+    @abstractmethod
+    def get_cut(self) -> str | None:
+        """
+        Return a string representing the cuts to apply to the data.
+        """
+        pass
+
+    @abstractmethod
+    def convert_to_pandas(self, data: dict) -> pd.DataFrame:
+        """
+        Convert the loaded data from a dictionary format to a pandas DataFrame.
+        """
+        pass
+
+    def _resolve_branches(self, all_branches: list) -> list[str]:
+        """Internal method to resolve wildcard patterns."""
+        selected = []
+        for pattern in self.get_features():
+            matched = fnmatch_filter(all_branches, pattern)
+            if not matched:
+                warnings.warn(f"'{pattern}' did not match any branches.")
+            selected.extend(matched)
+        return sorted(list(set(selected)))
+
+    def _save_to_parquet(self, df: pd.DataFrame, output_path: str):
+        """
+        Save the processed DataFrame to a file.
+        """
+        df.to_parquet(output_path)
+
+    def _save_to_csv(self, df: pd.DataFrame, output_path: str):
+        """
+        Save the processed DataFrame to a CSV file.
+        """
+        df.to_csv(output_path, index=False)
+
+    def process_file(self, file_path: str, out_file_path: str) -> pd.DataFrame:
+        """
+        Loads and processes a single ROOT file.
+        """
+
+        with uproot.open(f"{file_path}") as f:
+            tree = f[self.get_tree_name()]
+            all_branches = tree.keys()
+            branches_to_load = self._resolve_branches(all_branches)
+
+            if not branches_to_load:
+                return pd.DataFrame()
+
+            data = tree.arrays(branches_to_load, cut=self.get_cut(), how=dict)
+
+        df = self.convert_to_pandas(data)
+
+        if self.output_format == "parquet":
+            self._save_to_parquet(df, f"{out_file_path}.parquet")
+        elif self.output_format == "csv":
+            self._save_to_csv(df, f"{out_file_path}.csv")
+        else:
+            return pd.DataFrame()
+
+        return pd.DataFrame()

trigger_loader/cluster_manager.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+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__)
+
+
+class ClusterManager:
+    """Context manager to provision and tear down a Dask cluster.
+
+    Parameters
+    ----------
+    cluster_type : str
+        Backend to use ("local", "condor", "cuda", "kubernetes").
+    cluster_config : dict | None, optional
+        Keyword arguments forwarded to the specific cluster constructor.
+    jobs : int, optional
+        Desired number of jobs / workers (used for queue / scalable backends).
+    """
+
+    def __init__(
+        self,
+        cluster_type: str,
+        cluster_config: dict[str, Any] | None = None,
+        jobs: int = 1,
+    ) -> None:
+        if cluster_config is None:
+            cluster_config = {}
+        # Copy to avoid mutating caller's dict accidentally.
+        self.cluster_config: dict[str, Any] = dict(cluster_config)
+        self.cluster_type: str = cluster_type
+        self.jobs: int = jobs
+
+        self.cluster: Any | None = None
+        self.client: Any | None = None
+
+    # ------------------------------------------------------------------
+    # Context manager protocol
+    # ------------------------------------------------------------------
+    def __enter__(self):  # -> distributed.Client (avoids importing type eagerly)
+        self._start_cluster()
+        return self.client
+
+    def __exit__(self, exc_type, exc, tb) -> bool:  # noqa: D401 (simple)
+        self._close_cluster()
+        # Returning False propagates any exception (desired behavior)
+        return False
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    def _start_cluster(self) -> None:
+
+        ct = self.cluster_type.lower()
+
+        if ct == "local":
+            self.cluster = LocalCluster(**self.cluster_config)
+
+        elif ct == "condor":
+            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":
+            self.cluster = LocalCUDACluster(**self.cluster_config)
+
+        elif ct == "kubernetes":
+            self.cluster = KubeCluster(**self.cluster_config)
+            if self.jobs and self.jobs > 0:
+                try:
+                    # Not all KubeCluster versions expose scale() identically
+                    self.cluster.scale(self.jobs)
+                except Exception:
+                    pass  # Best effort; ignore if unsupported
+
+        else:
+            raise ValueError(f"Unsupported cluster type: {self.cluster_type}")
+
+        self.client = Client(self.cluster)
+        dash = getattr(self.client, "dashboard_link", None)
+        if dash:
+            logger.info(f"Dask dashboard: {dash}")
+
+    def _close_cluster(self) -> None:
+        # Close client first so tasks wind down before cluster termination.
+        if self.client is not None:
+            try:
+                self.client.close()
+            except Exception:
+                pass
+            finally:
+                self.client = None
+        if self.cluster is not None:
+            try:
+                self.cluster.close()
+            except Exception:
+                pass
+            finally:
+                self.cluster = None
+

trigger_loader/loader.py

@@ -0,0 +1,95 @@
+import json
+import logging
+import platform
+import time
+import uuid
+
+import awkward as ak
+import coffea
+from coffea import processor
+from coffea.nanoevents import NanoAODSchema
+
+from .cluster_manager import ClusterManager
+from .processor import TriggerProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class TriggerLoader:
+    def __init__(self,
+        sample_json: str,
+        transform: callable,
+        output_path: str,
+    ):
+        self.transform = transform
+        self.fileset = self._load_sample_json(sample_json)
+        self.output_path = output_path
+        self.run_uuid = str(uuid.uuid4())
+
+    def _build_processor(self):
+        run_meta = {
+            "run_uuid": self.run_uuid,
+            "fileset_size": sum(len(v) if isinstance(v, list) else 1 for v in self.fileset.values()),
+            "coffea_version": coffea.__version__,
+            "awkward_version": ak.__version__,
+            "python_version": platform.python_version(),
+        }
+
+        return TriggerProcessor(
+            output_path=self.output_path,
+            transform=self.transform,
+            compression="zstd",
+            add_uuid=False,
+            run_uuid=self.run_uuid,
+            run_metadata=run_meta,
+        )
+
+    def _load_sample_json(self, sample_json: str) -> dict:
+        with open(sample_json) as f:
+            return json.load(f)
+
+    def _write_run_metadata_file(self, path: str, duration_s: float | None = None):
+        meta_path = f"{path}/run_metadata.json"
+        data = {
+            "run_uuid": self.run_uuid,
+            "duration_seconds": duration_s,
+        }
+        with open(meta_path, "w") as f:
+            json.dump(data, f, indent=2)
+
+    def _run(self, runner: processor.Runner, label: str):
+        logger.log(f"Starting processing ({label})...")
+        start = time.time()
+        proc = self._build_processor()
+        acc = runner(
+            self.fileset,
+            treename="Events",
+            processor_instance=proc
+        )
+        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})")
+        return acc
+
+    def run_distributed(self, cluster_type: str, cluster_config: dict,
+                        chunksize: int = 100_000, jobs: int = 1):
+        with ClusterManager(cluster_type, cluster_config, jobs) as client:
+            executor = processor.DaskExecutor(client=client)
+            runner = processor.Runner(
+                executor=executor,
+                schema=NanoAODSchema,
+                chunksize=chunksize
+            )
+            self._run(runner, f"Distributed ({cluster_type})")
+
+    def run_local(self, num_workers: int = 4, chunksize: int = 100_000):
+        """
+        Run processing locally using a multi-processing executor.
+        """
+        executor = processor.FuturesExecutor(workers=num_workers)
+        runner = processor.Runner(
+            executor=executor,
+            schema=NanoAODSchema,
+            chunksize=chunksize
+        )
+        self._run(runner, f"Local ({num_workers} workers)")

trigger_loader/processor.py

@@ -0,0 +1,211 @@
+import datetime as dt
+import hashlib
+import json
+import os
+import time
+import uuid
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+import awkward as ak
+import pyarrow.parquet as pq
+from coffea import processor
+
+warnings.filterwarnings("ignore", message="Found duplicate branch")
+
+
+class TriggerProcessor(processor.ProcessorABC):
+    """
+    Coffea processor that applies a user transform to events and writes Parquet files.
+    
+    This processor transforms event data and writes the results to Parquet files with
+    comprehensive metadata tracking for reproducibility and provenance.
+    """
+
+    def __init__(
+        self,
+        output_path: str,
+        transform: Callable[[Any], ak.Array],
+        compression: str = "zstd",
+        compression_level: int | None = None,
+        filename_template: str = "{dataset}_{fileuuid}_{start}-{stop}.parquet",
+        add_uuid: bool = False,
+        write_manifest: bool = True,
+        manifest_name: str = "manifest.jsonl",
+        run_uuid: str | None = None,
+        run_metadata: dict | None = None,
+        preserve_event_metadata: bool = True,
+    ):
+        """
+        Initialize the TriggerProcessor.
+        
+        Args:
+            output_path: Directory where output files will be written
+            transform: Function to apply to events, returns awkward array
+            compression: Parquet compression algorithm
+            compression_level: Compression level (None for default)
+            filename_template: Template for output filenames
+            add_uuid: Whether to add UUID to filenames
+            write_manifest: Whether to write manifest file
+            manifest_name: Name of the manifest file
+            run_uuid: UUID for this processing run (generated if None)
+            run_metadata: Additional metadata for the run
+            preserve_event_metadata: Whether to preserve event-level metadata
+        """
+        self.output_path = output_path
+        self.transform = transform
+        self.compression = compression
+        self.compression_level = compression_level
+        self.filename_template = filename_template
+        self.add_uuid = add_uuid
+        self.write_manifest = write_manifest
+        self.manifest_name = manifest_name
+        self.run_uuid = run_uuid or str(uuid.uuid4())
+        self.run_metadata = run_metadata or {}
+        self.preserve_event_metadata = preserve_event_metadata
+
+        # Initialize output directory and paths
+        os.makedirs(self.output_path, exist_ok=True)
+        if write_manifest:
+            self._manifest_path = os.path.join(self.output_path, self.manifest_name)
+
+    @property
+    def accumulator(self):
+        """No aggregation needed (side-effect writing)."""
+        return {}
+
+    def process(self, events):
+        """
+        Process a chunk of events: transform and write to Parquet.
+        
+        Args:
+            events: Input events from Coffea
+            
+        Returns:
+            Empty dict (no accumulation needed)
+        """
+        # Apply transform and measure time
+        data, elapsed_s = self._apply_transform(events)
+
+        # Extract event metadata
+        event_meta = self._extract_event_metadata(events)
+        output_file = self._generate_output_filename(event_meta)
+
+        # Convert to Arrow table
+        table = ak.to_arrow_table(data)
+
+        # Build file metadata
+        file_meta = self._build_file_metadata(event_meta, table, elapsed_s)
+        table = self._embed_metadata_in_schema(table, file_meta, events.metadata)
+
+        # Write Parquet file
+        self._write_parquet(table, output_file)
+
+        # Write manifest entry
+        if self.write_manifest:
+            self._write_manifest_entry(output_file, file_meta)
+
+        return {}
+
+    def postprocess(self, accumulator):
+        """Postprocess accumulated results (no-op for this processor)."""
+        return accumulator
+
+    # ==================== Private Helper Methods ====================
+
+    def _apply_transform(self, events) -> tuple[ak.Array, float]:
+        """Apply user transform to events and measure execution time."""
+        t0 = time.time()
+        data = self.transform(events)
+        elapsed_s = time.time() - t0
+        return data, elapsed_s
+
+    def _extract_event_metadata(self, events) -> dict:
+        """Extract metadata from events object."""
+        source_file = None
+        if hasattr(events, "_events"):
+            source_file = getattr(events._events, "files", [None])[0]
+
+        return {
+            "start": events.metadata.get("entrystart", 0),
+            "stop": events.metadata.get("entrystop", 0),
+            "dataset": events.metadata.get("dataset", "unknown"),
+            "source_file": source_file,
+            "fileuuid": events.metadata.get("fileuuid"),
+        }
+
+    def _generate_output_filename(self, event_meta: dict) -> str:
+        """Generate output filename based on template and metadata."""
+        fname = self.filename_template.format(
+            dataset=event_meta["dataset"],
+            fileuuid=event_meta.get("fileuuid", "xx"),
+            start=event_meta["start"],
+            stop=event_meta["stop"]
+        )
+
+        if self.add_uuid:
+            stem, ext = os.path.splitext(fname)
+            fname = f"{stem}_{uuid.uuid4()}{ext}"
+
+        return os.path.join(self.output_path, fname)
+
+    def _build_file_metadata(self, event_meta: dict, table, elapsed_s: float) -> dict:
+        """Build comprehensive metadata dictionary for the output file."""
+        fileuuid = event_meta["fileuuid"]
+
+        return {
+            "run_uuid": self.run_uuid,
+            "dataset": event_meta["dataset"],
+            "source_root_file": event_meta["source_file"],
+            "fileuuid": str(fileuuid) if fileuuid is not None else None,
+            "entrystart": event_meta["start"],
+            "entrystop": event_meta["stop"],
+            "n_events_written": len(table),
+            "columns": table.column_names,
+            "created_utc": dt.datetime.now(dt.UTC).isoformat(timespec="seconds") + "Z",
+            "compression": self.compression,
+            "processing_time_s": round(elapsed_s, 6),
+        }
+
+    def _embed_metadata_in_schema(self, table, file_meta: dict, event_metadata: dict):
+        """Embed metadata into the Arrow table schema."""
+        schema = table.schema
+        existing = dict(schema.metadata or {})
+
+        # Add file metadata
+        existing[b"x-trigger-meta"] = json.dumps(
+            file_meta, separators=(",", ":")
+        ).encode()
+
+        # Optionally preserve event-level metadata
+        if self.preserve_event_metadata:
+            for k, v in event_metadata.items():
+                if k not in file_meta:
+                    file_meta[f"eventmeta_{k}"] = v
+
+        # Add run metadata hash
+        if self.run_metadata:
+            existing.setdefault(
+                b"x-run-meta-hash",
+                hashlib.sha256(
+                    json.dumps(self.run_metadata, sort_keys=True).encode()
+                ).hexdigest().encode()
+            )
+
+        return table.replace_schema_metadata(existing)
+
+    def _write_parquet(self, table, output_file: str):
+        """Write Arrow table to Parquet file."""
+        pq.write_table(
+            table,
+            output_file,
+            compression=self.compression,
+            compression_level=self.compression_level
+        )
+
+    def _write_manifest_entry(self, output_file: str, file_meta: dict):
+        """Write a single line to the manifest file."""
+        manifest_record = {"parquet_file": output_file, **file_meta}
+        with open(self._manifest_path, "a") as mf:
+            mf.write(json.dumps(manifest_record, separators=(",", ":")) + "\n")

triggerflow/cli.py

@@ -0,0 +1,122 @@
+import argparse
+import logging
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import yaml
+from cookiecutter.main import cookiecutter
+
+try:
+    from importlib.resources import as_file
+    from importlib.resources import files as ir_files
+except ImportError:
+    import importlib_resources
+    ir_files = importlib_resources.files  # type: ignore
+    as_file = importlib_resources.as_file  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+def echo(cmd): logger.info(f"$ {' '.join(str(c) for c in cmd)}", flush=True)
+def which_or_die(cmd):
+    if shutil.which(cmd) is None:
+        logger.info(f"Error: '{cmd}' not found on PATH.", file=sys.stderr); sys.exit(127)
+
+def packaged_starter_root() -> Path:
+    """Real FS path to packaged starter folder."""
+    base = ir_files("triggerflow")
+    with as_file(base / "starter") as p:
+        p = Path(p)
+        if p.exists():
+            return p
+    with as_file(base) as p:
+        logger.info("Error: starter not found. Contents of package:", file=sys.stderr)
+        for child in Path(p).iterdir(): logger.info(" -", child.name, file=sys.stderr)
+    sys.exit(2)
+
+def render_starter(project: str, out_dir: Path) -> Path:
+    """Render cookiecutter starter into out_dir; returns project path."""
+    starter = packaged_starter_root()
+    extra = {
+        "project_name": project,
+        "repo_name": project.replace("-", "_").lower(),
+        "python_package": project.replace("-", "_").lower(),
+    }
+
+    # cookiecutter.json expects bare keys, not cookiecutter.* here
+    out_dir.mkdir(parents=True, exist_ok=True)
+    proj_path = Path(cookiecutter(
+        template=str(starter),
+        no_input=True,
+        output_dir=str(out_dir),
+        extra_context=extra,
+    ))
+    return proj_path
+
+def find_env_yml(project_dir: Path) -> Path:
+    # Prefer root environment.yml; else find under src/**/
+    candidates = []
+    root = project_dir / "environment.yml"
+    if root.exists(): candidates.append(root)
+    candidates += list((project_dir / "src").rglob("environment.yml"))
+    if not candidates:
+        logger.info(f"Error: environment.yml not found under {project_dir}", file=sys.stderr); sys.exit(3)
+    # stable preference
+    candidates.sort(key=lambda p: (0 if p.parent.name != "src" else 1, len(str(p))))
+    return candidates[0]
+
+def conda_env_create_or_update(env_yml: Path) -> int:
+    which_or_die("conda")
+    # If YAML has 'name', override it to 'env_name' so updates are consistent
+    data = yaml.safe_load(env_yml.read_text(encoding="utf-8"))
+    if isinstance(data, dict):
+        tmp = env_yml.with_suffix(".rendered.yml")
+        tmp.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8")
+        env_yml = tmp
+
+    update = ["conda", "env", "update", "-f", str(env_yml), "--prune"]
+
+    echo(update); rc = subprocess.call(update)
+    if rc != 0: return rc
+
+    verify = ["conda", "run", "python", "-c", "import sys; logger.info(sys.executable)"]
+    echo(verify); subprocess.call(verify)
+    return 0
+
+def cmd_new(project: str, output: Path) -> int:
+    proj_dir = render_starter(project, output)
+    return 0
+
+def cmd_setup(project: str, output: Path) -> int:
+    # If project dir doesn’t exist yet, render it first
+    target = output / project
+    if not target.exists():
+        logger.info(f"Project '{project}' not found under {output}. Rendering starter first...")
+        render_starter(project, output)
+    env_yml = find_env_yml(target)
+    logger.info(f"Using environment file: {env_yml}")
+    return conda_env_create_or_update(env_yml)
+
+def main():
+    parser = argparse.ArgumentParser(prog="triggerflow", description="Triggerflow CLI")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    p_new = sub.add_parser("new", help="Render a new project from the packaged starter (Cookiecutter)")
+    p_new.add_argument("project", help="Project name")
+    p_new.add_argument("--out", default=".", help="Output directory (default: .)")
+
+    p_setup = sub.add_parser("setup", help="Create/update conda env from the rendered project's environment.yml")
+    p_setup.add_argument("project", help="Project/env name")
+    p_setup.add_argument("--out", default=".", help="Project parent directory (default: .)")
+
+    args = parser.parse_args()
+    out = Path(getattr(args, "out", ".")).resolve()
+    if args.cmd == "new":
+        sys.exit(cmd_new(args.project, out))
+    else:  # setup
+        sys.exit(cmd_setup(args.project, out))
+
+if __name__ == "__main__":
+    main()