making job submission and scripts more general, removing hard coded paths

root_gnn_dgl/jobs/cpu.sh

@@ -1 +1,5 @@
-salloc --nodes=1 --ntasks=64 --cpus-per-task=1 --qos=interactive --time=04:00:00 --constraint=cpu --account=atlas
+#!/bin/bash
+
+ACCOUNT="${SLURM_ACCOUNT:-atlas}"
+
+salloc --nodes=1 --ntasks=64 --cpus-per-task=1 --qos=interactive --time=04:00:00 --constraint=cpu --account="$ACCOUNT"

root_gnn_dgl/jobs/inference/run_inference.py

@@ -1,214 +1,277 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import glob
 import os
+import re
 import subprocess
-from mpi4py import MPI
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+try:
+    from mpi4py import MPI
+except ImportError:
+    MPI = None
 
-# Define the number of GPUs per node
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
 GPUS_PER_NODE = 4
 
-# Define your files, configs, and branch namesfiles = [
-files = [
-    "merged_aMCPy8_tHjb125_CP_0_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_0_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_0_FullSim_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_0_FullSim_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_15_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_15_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_30_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_30_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_45_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_45_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_45_FullSim_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_45_FullSim_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_60_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_60_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_75_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_75_AF3_lep_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_90_AF3_had_scaled.root",
-    "merged_aMCPy8_tHjb125_CP_90_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_0_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_0_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_0_FullSim_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_0_FullSim_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_15_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_15_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_30_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_30_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_45_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_45_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_45_FullSim_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_45_FullSim_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_60_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_60_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_75_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_75_AF3_lep_scaled.root",
-    "merged_aMCPy8_ttH125_CP_90_AF3_had_scaled.root",
-    "merged_aMCPy8_ttH125_CP_90_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_0_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_0_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_0_FullSim_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_0_FullSim_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_15_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_15_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_30_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_30_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_45_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_45_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_45_FullSim_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_45_FullSim_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_60_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_60_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_75_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_75_AF3_lep_scaled.root",
-    "merged_aMCPy8_tWH125_CP_90_AF3_had_scaled.root",
-    "merged_aMCPy8_tWH125_CP_90_AF3_lep_scaled.root",
-    "merged_hadnti_data.root",
-    "merged_hadti_data.root",
-    "merged_lepnti_data.root",
-    "merged_lepti_data.root",
-    "merged_MGPy8_ttyy_combined_FullSim_had_scaled.root",
-    "merged_MGPy8_ttyy_combined_FullSim_lep_scaled.root",
-    "merged_PhPy8_4FS_A14NNPDF_bbH125_FullSim_had_scaled.root",
-    "merged_PhPy8_4FS_A14NNPDF_bbH125_FullSim_lep_scaled.root",
-    "merged_PhPy8_ggH125_NNLOPS_FullSim_had_scaled.root",
-    "merged_PhPy8_ggH125_NNLOPS_FullSim_lep_scaled.root",
-    "merged_PhPy8_ggZH125_FullSim_had_scaled.root",
-    "merged_PhPy8_ggZH125_FullSim_lep_scaled.root",
-    "merged_PhPy8_VBFH125_FullSim_had_scaled.root",
-    "merged_PhPy8_VBFH125_FullSim_lep_scaled.root",
-    "merged_PhPy8_WmH125J_MINLO_FullSim_had_scaled.root",
-    "merged_PhPy8_WmH125J_MINLO_FullSim_lep_scaled.root",
-    "merged_PhPy8_WpH125J_MINLO_FullSim_had_scaled.root",
-    "merged_PhPy8_WpH125J_MINLO_FullSim_lep_scaled.root",
-    "merged_PhPy8_ZH125J_MINLO_FullSim_had_scaled.root",
-    "merged_PhPy8_ZH125J_MINLO_FullSim_lep_scaled.root",
-]
-# Configurations for each type of file
-config = [
-    "configs/model_configs/run_3_ttH/v06/sb_yukawa_no_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/sb_yukawa_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/tHjb_CP_no_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/tHjb_CP_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/ttH_CP_no_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/ttH_CP_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/ttH_tHjb_tWH_CP_no_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/ttH_tHjb_tWH_CP_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/tWH_CP_no_weights.yaml",
-    "configs/model_configs/run_3_ttH/v06/tWH_CP_weights.yaml",
-
-]
-
-# Branch names for each type of file
-branch_name = [
-    "sb_yukawa_no_weights_score",
-    "sb_yukawa_weights_score",
-    "tHjb_CP_no_weights_score",
-    "tHjb_CP_weights_score",
-    "ttH_CP_no_weights_score",
-    "ttH_CP_weights_score",
-    "ttH_tHjb_tWH_CP_no_weights_score",
-    "ttH_tHjb_tWH_CP_weights_score",
-    "tWH_CP_no_weights_score",
-    "tWH_CP_weights_score",
-]
-
-# MPI Initialization
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()  # Process rank (unique ID for each process)
-size = comm.Get_size()  # Total number of processes across all nodes
-
-# Master process manages the task queue
-def master():
-    num_workers = size - 1  # Exclude the master process
-    task_queue = files[:]  # Copy the list of files
-    workers_busy = [False] * num_workers  # Track if workers are busy
-
-    print(f"Master process starting with {num_workers} workers.")
-
-    # Distribute initial tasks to workers
-    for worker_rank in range(1, size):
-        if task_queue:
-            # Send a file to the worker
-            file = task_queue.pop(0)
-            comm.send(file, dest=worker_rank, tag=0)
-            workers_busy[worker_rank - 1] = True
-        else:
-            # No tasks left, send termination signal
-            comm.send(None, dest=worker_rank, tag=1)
-
-    # Process results and assign new tasks dynamically
-    while any(workers_busy):
-        # Receive results from workers
-        status = MPI.Status()
-        result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
-        worker_rank = status.Get_source()
-        workers_busy[worker_rank - 1] = False  # Worker is now free
-
-        # Assign a new task if available
-        if task_queue:
-            file = task_queue.pop(0)
-            comm.send(file, dest=worker_rank, tag=0)
-            workers_busy[worker_rank - 1] = True
+
+@dataclass(frozen=True)
+class InferenceTask:
+    target: Path
+    destination: Path
+
+
+def load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def repo_relative(path: Path) -> str:
+    try:
+        return str(path.resolve().relative_to(REPO_ROOT))
+    except ValueError:
+        return str(path)
+
+
+def resolve_path(raw_path: str) -> Path:
+    path = Path(raw_path).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+
+    repo_path = (REPO_ROOT / path).resolve()
+    if repo_path.exists() or str(raw_path).startswith(("configs/", "jobs/", "scripts/")):
+        return repo_path
+
+    return (Path.cwd() / path).resolve()
+
+
+def sanitize_branch_name(name: str) -> str:
+    cleaned = re.sub(r"\W+", "_", name).strip("_")
+    if not cleaned:
+        cleaned = "gnn"
+    if cleaned[0].isdigit():
+        cleaned = f"gnn_{cleaned}"
+    return f"{cleaned}_score"
+
+
+def discover_model_configs(raw_configs: list[str], raw_config_dirs: list[str]) -> list[Path]:
+    configs = []
+    for raw_config in raw_configs:
+        config = resolve_path(raw_config)
+        if not config.is_file():
+            raise FileNotFoundError(f"Config file does not exist: {config}")
+        configs.append(config)
+
+    for raw_dir in raw_config_dirs:
+        config_dir = resolve_path(raw_dir)
+        if not config_dir.is_dir():
+            raise NotADirectoryError(f"Config directory does not exist: {config_dir}")
+        configs.extend(sorted(path.resolve() for path in config_dir.glob("*.yaml")))
+
+    unique_configs = []
+    seen = set()
+    for config in configs:
+        if config not in seen:
+            unique_configs.append(config)
+            seen.add(config)
+    if not unique_configs:
+        raise ValueError("No model configs were provided.")
+    return unique_configs
+
+
+def default_branch_names(config_paths: list[Path]) -> list[str]:
+    branch_names = []
+    for config_path in config_paths:
+        config = load_config(config_path)
+        training_name = config.get("Training_Name") or config_path.stem
+        branch_names.append(sanitize_branch_name(str(training_name)))
+    return branch_names
+
+
+def discover_targets_from_config(sample_config_path: Path) -> list[Path]:
+    config = load_config(sample_config_path)
+    targets = []
+
+    for dataset_name, dataset_config in config.get("Datasets", {}).items():
+        args = dataset_config.get("args", {})
+        raw_dir = args.get("raw_dir")
+        file_names = args.get("file_names")
+        if not raw_dir or not file_names:
+            raise ValueError(f"Dataset {dataset_name} is missing raw_dir or file_names.")
+
+        patterns = file_names if isinstance(file_names, list) else [file_names]
+        for pattern in patterns:
+            matches = sorted(Path(path).resolve() for path in glob.glob(os.path.join(raw_dir, pattern)))
+            if not matches:
+                raise FileNotFoundError(f"No files matched {os.path.join(raw_dir, pattern)}")
+            targets.extend(matches)
+
+    unique_targets = []
+    seen = set()
+    for target in targets:
+        if target not in seen:
+            unique_targets.append(target)
+            seen.add(target)
+    return unique_targets
+
+
+def discover_targets(raw_targets: list[str], sample_config: str | None) -> list[Path]:
+    targets = []
+    if sample_config:
+        targets.extend(discover_targets_from_config(resolve_path(sample_config)))
+
+    for raw_target in raw_targets:
+        matches = sorted(Path(path).resolve() for path in glob.glob(str(resolve_path(raw_target))))
+        if matches:
+            targets.extend(matches)
         else:
-            # No tasks left, send termination signal
-            comm.send(None, dest=worker_rank, tag=1)
-
-    print("Master process: All tasks completed.")
-
-# Worker processes handle tasks dynamically
-def worker():
-    gpu_id = (rank - 1) % GPUS_PER_NODE  # Assign GPU ID based on rank
-    print(f"Worker {rank} starting, GPU ID: {gpu_id}", flush=True)
-
-    while True:
-        # Receive a task from the master
-        file = comm.recv(source=0, tag=MPI.ANY_TAG)
-        if file is None:
-            # Termination signal received
-            print(f"Worker {rank}: Termination signal received.")
-            break
-
-        # Initialize config and branch_name for each file
-        local_config = config[:]  # Create a copy of the global config
-        local_branch_name = branch_name[:]  # Create a copy of the global branch_name
-
-        # Determine file type and append appropriate configurations
-        if "had" in file:
-            local_config += had_config  # Append had-specific configurations
-            local_branch_name += had_branch_name  # Append had-specific branch names
-        elif "lep" in file:
-            local_config += lep_config  # Append lep-specific configurations
-            local_branch_name += lep_branch_name  # Append lep-specific branch names
-
-        # Process the file
-        print(f"Worker {rank}: Processing file {file} on GPU {gpu_id}", flush=True)
-        print(f"branches = {local_branch_name}")
-        target = f"/global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/v02/preselection/merged/{file}"
-        destination = f"/global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/scores/gnn_v06/{file}"
-        
-        # Set the GPU for this process
-        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
-        
-        # Build the command
-        command = [
-            "python", "scripts/inference.py",
-            "--target", target,
-            "--destination", destination,
-            "--config", *local_config,
-            "--branch_name", *local_branch_name,
-            "--chunks", "1",
-            "--chunkno", "0",
-            "--write"
-        ]
-        
-        # Execute the command
-        subprocess.run(command)
-        
-        # Notify the master that the task is complete
-        comm.send(f"Worker {rank}: Finished {file}", dest=0, tag=0)
-
-# Main function
-if __name__ == "__main__":
+            target = resolve_path(raw_target)
+            if not target.exists():
+                raise FileNotFoundError(f"Target does not exist: {target}")
+            targets.append(target)
+
+    unique_targets = []
+    seen = set()
+    for target in targets:
+        if target not in seen:
+            unique_targets.append(target)
+            seen.add(target)
+    if not unique_targets:
+        raise ValueError("No inference targets were provided.")
+    return unique_targets
+
+
+def build_tasks(targets: list[Path], output_dir: Path) -> list[InferenceTask]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    return [InferenceTask(target=target, destination=(output_dir / target.name).resolve()) for target in targets]
+
+
+def build_command(
+    task: InferenceTask,
+    config_paths: list[Path],
+    branch_names: list[str],
+    args: argparse.Namespace,
+) -> list[str]:
+    command = [
+        sys.executable,
+        "scripts/inference.py",
+        "--target",
+        str(task.target),
+        "--destination",
+        str(task.destination),
+        "--config",
+        *[str(path) for path in config_paths],
+        "--branch_name",
+        *branch_names,
+        "--chunks",
+        str(args.chunks),
+        "--chunkno",
+        str(args.chunkno),
+    ]
+
+    if args.write:
+        command.append("--write")
+    if args.clobber:
+        command.append("--clobber")
+    if args.ckpt is not None:
+        command.extend(["--ckpt", str(args.ckpt)])
+    if args.var:
+        command.extend(["--var", args.var])
+    if args.mode:
+        command.extend(["--mode", args.mode])
+    if args.tree:
+        command.extend(["--tree", args.tree])
+    return command
+
+
+def run_task(
+    task: InferenceTask,
+    config_paths: list[Path],
+    branch_names: list[str],
+    args: argparse.Namespace,
+    rank: int,
+) -> int:
+    gpu_id = rank % GPUS_PER_NODE
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+    command = build_command(task, config_paths, branch_names, args)
+
+    print(
+        f"[rank {rank}] target={task.target} destination={task.destination} gpu={gpu_id}",
+        flush=True,
+    )
+    print(f"[rank {rank}] command={' '.join(command)}", flush=True)
+    if args.test:
+        return 0
+    return subprocess.run(command, cwd=REPO_ROOT, env=env).returncode
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run scripts/inference.py over many ROOT files with many model configs."
+    )
+    parser.add_argument("--sample-config", help="Config whose dataset file_names define inference targets.")
+    parser.add_argument("--target", nargs="*", default=[], help="Target ROOT files or glob patterns.")
+    parser.add_argument("--config", nargs="*", default=[], help="Model config files to score.")
+    parser.add_argument("--config-dir", nargs="*", default=[], help="Directories containing model YAML configs.")
+    parser.add_argument("--branch-name", nargs="*", default=[], help="Score branch names. Defaults to Training_Name_score.")
+    parser.add_argument("--output-dir", required=True, help="Directory for output ROOT/NPZ files.")
+    parser.add_argument("--chunks", type=int, default=1)
+    parser.add_argument("--chunkno", type=int, default=0)
+    parser.add_argument("--write", action="store_true", help="Write ROOT files with score branches.")
+    parser.add_argument("--clobber", action="store_true")
+    parser.add_argument("--ckpt", type=int, default=None, help="Checkpoint epoch. Omit to let scripts/inference.py choose best epoch.")
+    parser.add_argument("--var", default="Test_AUC")
+    parser.add_argument("--mode", default="max")
+    parser.add_argument("--tree", default="")
+    parser.add_argument("--test", action="store_true", help="Print planned commands without running inference.")
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    config_paths = discover_model_configs(args.config, args.config_dir)
+    branch_names = args.branch_name or default_branch_names(config_paths)
+    if len(branch_names) != len(config_paths):
+        raise ValueError("--branch-name count must match the number of model configs.")
+
+    targets = discover_targets(args.target, args.sample_config)
+    tasks = build_tasks(targets, resolve_path(args.output_dir))
+
+    print(f"Repo root: {REPO_ROOT}", flush=True)
+    print(f"Targets: {len(tasks)}", flush=True)
+    print(f"Model configs: {len(config_paths)}", flush=True)
+    print("Branches:", " ".join(branch_names), flush=True)
+
+    if MPI is None:
+        failures = 0
+        for task in tasks:
+            failures += int(run_task(task, config_paths, branch_names, args, rank=0) != 0)
+        return 1 if failures else 0
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    failures = 0
+    for index, task in enumerate(tasks):
+        if index % size != rank:
+            continue
+        failures += int(run_task(task, config_paths, branch_names, args, rank=rank) != 0)
+
+    total_failures = comm.allreduce(failures, op=MPI.SUM)
     if rank == 0:
-        master()
-    else:
-        worker()
+        print(f"Completed {len(tasks) - total_failures} task(s); failed {total_failures}.", flush=True)
+    return 1 if total_failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

root_gnn_dgl/jobs/interactive.sh

@@ -1 +1,5 @@
-salloc --nodes 1 --qos shared_interactive --time 04:00:00 --constraint gpu --account=atlas --gres=gpu:1
+#!/bin/bash
+
+ACCOUNT="${SLURM_ACCOUNT:-atlas}"
+
+salloc --nodes 1 --qos shared_interactive --time 04:00:00 --constraint gpu --account="$ACCOUNT" --gres=gpu:1

root_gnn_dgl/jobs/prep_data/run_processing.py

@@ -1,7 +1,9 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
+import argparse
 import yaml
 import os
 import subprocess
+from pathlib import Path
 
 def extract_dataset_keys_and_chunks(yaml_file):
     """
@@ -34,12 +36,11 @@ def call_bash_script(config, dataset_key, shuffle_chunks, bash_script):
     """
     try:
         # Construct the command to call the bash script
-        command = f"{bash_script} {config} {dataset_key} {shuffle_chunks}"
+        command = [bash_script, config, dataset_key, str(shuffle_chunks)]
 
-        print(f"Executing command: {command}")  # Debugging: Print the command being executed
+        print(f"Executing command: {' '.join(command)}")
 
-        # Run the command
-        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
+        result = subprocess.run(command, check=True, capture_output=True, text=True)
         if result.stdout:
             print(result.stdout, end="")
         if result.stderr:
@@ -55,7 +56,7 @@ def process_yaml_file(config, base_directory, bash_script):
     """
     Process a single YAML file by extracting datasets and calling the bash script in parallel.
     """
-    yaml_file = os.path.join(base_directory, config)  # Construct full path
+    yaml_file = os.path.join(base_directory, config)
     if os.path.exists(yaml_file):
         print(f"Processing file: {config}")
         datasets = extract_dataset_keys_and_chunks(yaml_file)
@@ -77,18 +78,19 @@ def process_yaml_file(config, base_directory, bash_script):
         return [f"File not found: {yaml_file}"]
 
 def main():
-    # Base directory containing the YAML files
-    base_directory = os.getcwd() + "/"
+    parser = argparse.ArgumentParser(description="Run prep_data.sh for every dataset in one or more configs.")
+    parser.add_argument("configs", nargs="+", help="YAML config files, relative to the repo root or absolute.")
+    parser.add_argument(
+        "--bash-script",
+        default=None,
+        help="Path to prep_data.sh. Defaults to jobs/prep_data/prep_data.sh under the repo root.",
+    )
+    args = parser.parse_args()
 
-    # List of YAML configuration files (relative to the base directory)
-    configs = [
-        # "configs/stats_100K/pretraining_multiclass.yaml",
-        "configs/stats_100K/ttH_CP_even_vs_odd.yaml",
-        # "configs/stats_all/pretraining_multiclass.yaml",
-    ]
-
-    # Path to the bash script to be called
-    bash_script = base_directory + "jobs/prep_data/prep_data.sh"
+    repo_root = Path(__file__).resolve().parents[2]
+    base_directory = str(repo_root) + "/"
+    configs = [str(Path(config).resolve()) if Path(config).is_absolute() else config for config in args.configs]
+    bash_script = args.bash_script or str(repo_root / "jobs/prep_data/prep_data.sh")
 
     # Use ThreadPoolExecutor to process YAML files concurrently
     max_workers = os.cpu_count()  # Use all available CPU cores

root_gnn_dgl/jobs/salloc.sh

@@ -1 +1,5 @@
-salloc --nodes 4 --qos interactive --time 04:00:00 --constraint gpu --account=atlas --gres=gpu:4
+#!/bin/bash
+
+ACCOUNT="${SLURM_ACCOUNT:-atlas}"
+
+salloc --nodes 4 --qos interactive --time 04:00:00 --constraint gpu --account="$ACCOUNT" --gres=gpu:4

root_gnn_dgl/jobs/training/conda/run_job.sh

@@ -3,26 +3,32 @@
 #SBATCH -C gpu
 #SBATCH -q shared
 #SBATCH -t 15:00:00
-#SBATCH -A atlas
-#SBATCH -o /global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/jobs/slurm/%j.out # STDOUT
+#SBATCH -o jobs/slurm/%j.out # STDOUT
 
 CONFIG=$1
 shift
-ARGUEMENTS="$*"
+ARGUMENTS=("$@")
 
-DIRECTORY="/global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/"
-BASE_COMMAND="$DIRECTORY/training_script.py $ARGUEMENTS --preshuffle --nocompile --lazy --config $DIRECTORY"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
 
-echo "launched image"
-cd $DIRECTORY
+cd "$REPO_ROOT"
+mkdir -p jobs/slurm
 
 eval "$(conda shell.bash hook)"
-conda init bash
-conda activate /opt/conda/envs/dgl
+conda activate "${ROOT_GNN_CONDA_ENV:-dgl}"
 
-COMMAND="$BASE_COMMAND$CONFIG"
+COMMAND=(
+    "$REPO_ROOT/scripts/training_script.py"
+    "${ARGUMENTS[@]}"
+    "--preshuffle"
+    "--nocompile"
+    "--lazy"
+    "--config"
+    "$CONFIG"
+)
 
 echo "Running my script now"
-echo $COMMAND
-python -u $COMMAND
-echo "Done"
+echo "Executing: python -u ${COMMAND[*]}"
+python -u "${COMMAND[@]}"
+echo "Done"

root_gnn_dgl/jobs/training/conda/submit.sh

@@ -1,14 +1,28 @@
+#!/bin/bash
+
 date
 
-DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+RUN_JOB="$REPO_ROOT/jobs/training/conda/run_job.sh"
+SBATCH_ACCOUNT_ARGS=()
+if [ -n "${SLURM_ACCOUNT:-}" ]; then
+  SBATCH_ACCOUNT_ARGS=(--account "$SLURM_ACCOUNT")
+fi
 
-configs=(  
+default_configs=(
   "configs/stats_all/ttH_CP_even_vs_odd.yaml"
   "configs/stats_all/ttH_CP_even_vs_odd_batch_size_2048.yaml"
   "configs/stats_all/ttH_CP_even_vs_odd_batch_size_4096.yaml"
   "configs/stats_all/ttH_CP_even_vs_odd_batch_size_8192.yaml"
 )
 
+if [ "$#" -gt 0 ]; then
+  configs=("$@")
+else
+  configs=("${default_configs[@]}")
+fi
+
 counter=0
 
 hours=12
@@ -16,8 +30,8 @@ time="${hours}:00:00"
 
 for job in "${configs[@]}"
 do
-  sbatch --job-name="$job" --time="$time" /global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/jobs/training/conda/run_job.sh "$job"
+  sbatch "${SBATCH_ACCOUNT_ARGS[@]}" --job-name="$job" --time="$time" "$RUN_JOB" "$job"
   ((counter++))
 done
 
-echo "Total jobs submitted: $counter"
+echo "Total jobs submitted: $counter"

root_gnn_dgl/jobs/training/podman/run_job.sh

@@ -3,12 +3,16 @@
 #SBATCH -C "gpu&hbm80g"
 #SBATCH -q shared
 #SBATCH -t 24:00:00
-#SBATCH -A atlas
-#SBATCH -o /global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/jobs/slurm/%j.out # STDOUT
+#SBATCH -o jobs/slurm/%j.out # STDOUT
 
-ARGUEMENTS="$*"
+ARGUMENTS=("$@")
 
-echo "Arguements: $ARGUEMENTS"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
 
-## create a launch image script
-source "/global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/setup/launch_image.sh" "/global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/jobs/training/podman/run_job_image.sh" $ARGUEMENTS
+cd "$REPO_ROOT"
+mkdir -p jobs/slurm
+
+echo "Arguments: ${ARGUMENTS[*]}"
+
+source "$REPO_ROOT/setup/launch_image.sh" "$REPO_ROOT/jobs/training/podman/run_job_image.sh" "${ARGUMENTS[@]}"

root_gnn_dgl/jobs/training/podman/run_job_image.sh

@@ -5,20 +5,21 @@ shift
 # Store any other potential arguments safely
 OTHER_ARGUEMENTS=("$@")
 
-DIRECTORY="/global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
 
-cd $DIRECTORY
+cd "$REPO_ROOT"
 
 # Use a bash array to build the command and its arguments
 # Each element in the () is a separate argument.
 COMMAND_ARGS=(
-    "$DIRECTORY/scripts/training_script.py"
+    "$REPO_ROOT/scripts/training_script.py"
     "${OTHER_ARGUEMENTS[@]}"
     "--preshuffle"
     "--nocompile"
     "--lazy"
     "--config"
-    "$DIRECTORY$CONFIG"
+    "$CONFIG"
 )
 
 echo "Running my script now"
@@ -28,4 +29,4 @@ echo "Executing: python -u ${COMMAND_ARGS[@]}"
 # The "${COMMAND_ARGS[@]}" syntax ensures each element is passed as a distinct argument
 python -u "${COMMAND_ARGS[@]}"
 
-echo "Done"
+echo "Done"

root_gnn_dgl/jobs/training/podman/submit.sh

@@ -1,8 +1,16 @@
+#!/bin/bash
+
 date
 
-DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+RUN_JOB="$REPO_ROOT/jobs/training/podman/run_job.sh"
+SBATCH_ACCOUNT_ARGS=()
+if [ -n "${SLURM_ACCOUNT:-}" ]; then
+  SBATCH_ACCOUNT_ARGS=(--account "$SLURM_ACCOUNT")
+fi
 
-configs=(  
+default_configs=(
   # "configs/multiclass_pretraining/baseline.yaml"
   # "configs/multiclass_pretraining/pretraining_batch_size/multiclass_bs_4096.yaml"
   # "configs/multiclass_pretraining/pretraining_hid_size/multiclass_hid_256.yaml"
@@ -47,6 +55,12 @@ configs=(
 
 )
 
+if [ "$#" -gt 0 ]; then
+  configs=("$@")
+else
+  configs=("${default_configs[@]}")
+fi
+
 counter=0
 
 hours=24
@@ -54,7 +68,7 @@ time="${hours}:00:00"
 
 for job in "${configs[@]}"
 do
-  sbatch --job-name="$job" --time="$time" /global/cfs/projectdirs/atlas/joshua/gnn/root_gnn_dgl/jobs/training/podman/run_job.sh "$job"
+  sbatch "${SBATCH_ACCOUNT_ARGS[@]}" --job-name="$job" --time="$time" "$RUN_JOB" "$job"
   ((counter++))
 done
 

root_gnn_dgl/jobs/training/run_parallel_trainings.py

@@ -0,0 +1,509 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+import shlex
+import subprocess
+import sys
+import time
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+import socket
+
+import yaml
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_TARGETS = (
+    REPO_ROOT / "configs/higgs_production_updated_leptons",
+    REPO_ROOT / "configs/triboson",
+)
+FIXED_GPU_TOKENS = ("0", "1", "2", "3")
+DEFAULT_TRAINING_ARGS = ("--preshuffle", "--nocompile", "--lazy")
+FORBIDDEN_TRAINING_ARGS = {"--config", "--multigpu", "--multinode"}
+
+
+@dataclass(frozen=True)
+class ExecutionSlot:
+    gpu_token: str
+
+
+@dataclass
+class TrainingJob:
+    config_path: Path
+    training_directory: str
+    pretraining_path: Path | None
+
+
+@dataclass
+class ActiveProcess:
+    job: TrainingJob
+    slot: ExecutionSlot
+    log_path: Path
+    log_handle: object
+    process: subprocess.Popen
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Queue local single-GPU trainings on the current node only. This launcher "
+            "is intentionally hard-coded to use GPUs 0,1,2,3 with no srun or multi-node logic."
+        ),
+        epilog=(
+            "Unknown arguments are forwarded to scripts/training_script.py. "
+            "Example: python jobs/training/run_parallel_trainings.py --split 1/2 --restart --seed 7"
+        ),
+    )
+    parser.add_argument(
+        "targets",
+        nargs="*",
+        help=(
+            "Config files or directories. Defaults to "
+            "configs/higgs_production_updated_leptons and configs/triboson."
+        ),
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        help=(
+            "Run only one deterministic slice of the config queue, formatted as K/N. "
+            "Example: --split 1/2 for the first half and --split 2/2 for the second."
+        ),
+    )
+    parser.add_argument(
+        "--logs-root",
+        type=str,
+        default="jobs/slurm/parallel_training_logs",
+        help="Directory where launcher stdout/stderr logs should be written.",
+    )
+    parser.add_argument(
+        "--stop-on-failure",
+        action="store_true",
+        help="Stop launching new trainings after the first failure.",
+    )
+    parser.add_argument(
+        "--ignore-missing-pretraining",
+        action="store_true",
+        help="Launch even if a finetuning config points to a missing pretraining checkpoint.",
+    )
+    parser.add_argument(
+        "--test",
+        action="store_true",
+        help="Run preflight checks and print the launch plan without starting trainings.",
+    )
+    parser.add_argument(
+        "--poll-seconds",
+        type=float,
+        default=10.0,
+        help="Seconds between checks for finished trainings.",
+    )
+
+    # Accepted only so older commands do not fail; these are ignored in the
+    # single-node hard-coded launcher.
+    parser.add_argument("--gpus", type=str, help=argparse.SUPPRESS)
+    parser.add_argument("--nodelist", type=str, help=argparse.SUPPRESS)
+    parser.add_argument("--local-only", action="store_true", help=argparse.SUPPRESS)
+    parser.add_argument("--srun-binary", type=str, help=argparse.SUPPRESS)
+    parser.add_argument("--scontrol-binary", type=str, help=argparse.SUPPRESS)
+    parser.add_argument("--max-parallel", type=int, help=argparse.SUPPRESS)
+    parser.add_argument("--dry-run", action="store_true", help=argparse.SUPPRESS)
+
+    args, forwarded_args = parser.parse_known_args()
+    return args, forwarded_args
+
+
+def repo_relative(path: Path) -> str:
+    try:
+        return str(path.resolve().relative_to(REPO_ROOT))
+    except ValueError:
+        return str(path.resolve())
+
+
+def resolve_target(raw_target: str) -> Path:
+    raw_path = Path(raw_target).expanduser()
+    if raw_path.is_absolute():
+        return raw_path.resolve()
+
+    repo_candidate = (REPO_ROOT / raw_path).resolve()
+    if repo_candidate.exists():
+        return repo_candidate
+
+    return (Path.cwd() / raw_path).resolve()
+
+
+def discover_configs(raw_targets: list[str]) -> list[Path]:
+    targets = [resolve_target(raw_target) for raw_target in raw_targets] if raw_targets else list(DEFAULT_TARGETS)
+    configs: list[Path] = []
+
+    for target in targets:
+        if target.is_file():
+            if target.suffix != ".yaml":
+                raise FileNotFoundError(f"Expected a .yaml config file, got {target}")
+            configs.append(target)
+            continue
+
+        if not target.is_dir():
+            raise FileNotFoundError(f"Config target does not exist: {target}")
+
+        configs.extend(sorted(path.resolve() for path in target.glob("*.yaml")))
+
+    unique_configs = []
+    seen = set()
+    for config in configs:
+        if config not in seen:
+            unique_configs.append(config)
+            seen.add(config)
+    return unique_configs
+
+
+def resolve_checkpoint_path(raw_path: str) -> Path:
+    candidate = Path(raw_path).expanduser()
+    if candidate.is_absolute():
+        return candidate.resolve()
+    return (REPO_ROOT / candidate).resolve()
+
+
+def load_job(config_path: Path) -> TrainingJob:
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+
+    if not isinstance(config, dict):
+        raise ValueError(f"Config is not a YAML mapping: {config_path}")
+
+    training_directory = config.get("Training_Directory")
+    if not training_directory:
+        raise ValueError(f"Config is missing Training_Directory: {config_path}")
+
+    pretraining_path = None
+    model_args = config.get("Model", {}).get("args", {})
+    if isinstance(model_args, dict) and model_args.get("pretraining_path"):
+        pretraining_path = resolve_checkpoint_path(model_args["pretraining_path"])
+
+    return TrainingJob(
+        config_path=config_path.resolve(),
+        training_directory=str(training_directory),
+        pretraining_path=pretraining_path,
+    )
+
+
+def parse_split_spec(raw_split: str | None) -> tuple[int, int] | None:
+    if raw_split is None:
+        return None
+
+    match = re.fullmatch(r"\s*(\d+)\s*/\s*(\d+)\s*", raw_split)
+    if not match:
+        raise ValueError(
+            f"Invalid --split value {raw_split!r}. Use the format K/N, for example 1/2."
+        )
+
+    split_index = int(match.group(1))
+    split_count = int(match.group(2))
+
+    if split_count <= 0:
+        raise ValueError("--split requires N > 0.")
+    if split_index <= 0:
+        raise ValueError("--split requires K >= 1.")
+    if split_index > split_count:
+        raise ValueError("--split requires K <= N.")
+
+    return split_index, split_count
+
+
+def select_jobs_for_split(
+    jobs: list[TrainingJob],
+    split_spec: tuple[int, int] | None,
+) -> tuple[list[TrainingJob], str | None]:
+    if split_spec is None:
+        return jobs, None
+
+    split_index, split_count = split_spec
+    selected_jobs = [
+        job for index, job in enumerate(jobs)
+        if index % split_count == (split_index - 1)
+    ]
+    return selected_jobs, f"{split_index}/{split_count}"
+
+
+def validate_forwarded_args(forwarded_args: list[str]) -> None:
+    normalized_args = [arg.split("=", 1)[0] for arg in forwarded_args]
+    forbidden = sorted(arg for arg in normalized_args if arg in FORBIDDEN_TRAINING_ARGS)
+    if forbidden:
+        joined = ", ".join(forbidden)
+        raise ValueError(
+            f"These training arguments are controlled by the launcher and cannot be forwarded: {joined}"
+        )
+
+
+def warn_about_ignored_args(args: argparse.Namespace) -> None:
+    ignored = []
+
+    if args.gpus:
+        ignored.append("--gpus")
+    if args.nodelist:
+        ignored.append("--nodelist")
+    if args.local_only:
+        ignored.append("--local-only")
+    if args.srun_binary:
+        ignored.append("--srun-binary")
+    if args.scontrol_binary:
+        ignored.append("--scontrol-binary")
+    if args.max_parallel is not None:
+        ignored.append("--max-parallel")
+    if args.dry_run:
+        ignored.append("--dry-run")
+
+    if ignored:
+        print(
+            "Ignoring launcher options not used by the single-node version: "
+            + ", ".join(ignored),
+            flush=True,
+        )
+
+
+def format_slot(slot: ExecutionSlot, hostname: str) -> str:
+    return f"{hostname}:gpu{slot.gpu_token}"
+
+
+def sanitize_log_name(config_path: Path) -> str:
+    try:
+        relative = config_path.resolve().relative_to(REPO_ROOT)
+        name = "__".join(relative.parts)
+    except ValueError:
+        name = config_path.name
+    return f"{name}.log"
+
+
+def terminate_active_processes(active_processes: list[ActiveProcess]) -> None:
+    for active in active_processes:
+        if active.process.poll() is None:
+            active.process.terminate()
+
+    deadline = time.time() + 10
+    while time.time() < deadline:
+        if all(active.process.poll() is not None for active in active_processes):
+            break
+        time.sleep(0.5)
+
+    for active in active_processes:
+        if active.process.poll() is None:
+            active.process.kill()
+
+    for active in active_processes:
+        if not active.log_handle.closed:
+            active.log_handle.close()
+
+
+def build_training_command(job: TrainingJob, forwarded_args: list[str]) -> list[str]:
+    return [
+        sys.executable,
+        "scripts/training_script.py",
+        "--config",
+        str(job.config_path),
+        *DEFAULT_TRAINING_ARGS,
+        *forwarded_args,
+    ]
+
+
+def launch_job(
+    job: TrainingJob,
+    slot: ExecutionSlot,
+    forwarded_args: list[str],
+    log_dir: Path,
+    hostname: str,
+) -> ActiveProcess:
+    log_path = log_dir / sanitize_log_name(job.config_path)
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    log_handle = log_path.open("w", encoding="utf-8")
+
+    command = build_training_command(job, forwarded_args)
+
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = slot.gpu_token
+
+    process = subprocess.Popen(
+        command,
+        cwd=REPO_ROOT,
+        env=env,
+        stdout=log_handle,
+        stderr=subprocess.STDOUT,
+    )
+
+    print(
+        f"[launch] slot={format_slot(slot, hostname)} config={repo_relative(job.config_path)} "
+        f"log={repo_relative(log_path)}",
+        flush=True,
+    )
+
+    return ActiveProcess(
+        job=job,
+        slot=slot,
+        log_path=log_path,
+        log_handle=log_handle,
+        process=process,
+    )
+
+
+def main() -> int:
+    args, forwarded_args = parse_args()
+    validate_forwarded_args(forwarded_args)
+    warn_about_ignored_args(args)
+    split_spec = parse_split_spec(args.split)
+
+    config_paths = discover_configs(args.targets)
+    if not config_paths:
+        raise RuntimeError("No YAML configs were found for the requested targets.")
+
+    all_jobs = [load_job(config_path) for config_path in config_paths]
+    jobs, split_label = select_jobs_for_split(all_jobs, split_spec)
+    if not jobs:
+        raise RuntimeError("No configs were selected for this split.")
+
+    duplicated_training_dirs: dict[str, Path] = {}
+    duplicate_messages = []
+    missing_pretraining = []
+    for job in jobs:
+        if job.training_directory in duplicated_training_dirs:
+            duplicate_messages.append(
+                f"{repo_relative(job.config_path)} conflicts with "
+                f"{repo_relative(duplicated_training_dirs[job.training_directory])} "
+                f"on Training_Directory={job.training_directory}"
+            )
+        else:
+            duplicated_training_dirs[job.training_directory] = job.config_path
+
+        if job.pretraining_path is not None and not job.pretraining_path.exists():
+            missing_pretraining.append(
+                f"{repo_relative(job.config_path)} -> missing {job.pretraining_path}"
+            )
+
+    if duplicate_messages:
+        raise RuntimeError(
+            "Refusing to launch configs with overlapping Training_Directory values:\n"
+            + "\n".join(duplicate_messages)
+        )
+
+    if missing_pretraining and not args.ignore_missing_pretraining:
+        raise RuntimeError(
+            "Missing pretraining checkpoints were found:\n"
+            + "\n".join(missing_pretraining)
+            + "\nRe-run with --ignore-missing-pretraining to launch anyway."
+        )
+
+    if missing_pretraining:
+        print("Pretraining checkpoint warnings:", flush=True)
+        for message in missing_pretraining:
+            print(f"  {message}", flush=True)
+
+    hostname = socket.gethostname().split(".", 1)[0]
+    execution_slots = [ExecutionSlot(gpu_token=gpu_token) for gpu_token in FIXED_GPU_TOKENS]
+    max_parallel = len(execution_slots)
+
+    timestamp = time.strftime("%Y%m%d_%H%M%S")
+    log_dir = (REPO_ROOT / args.logs_root / timestamp).resolve()
+
+    print(f"Repo root: {REPO_ROOT}", flush=True)
+    print(f"Configs discovered: {len(all_jobs)}", flush=True)
+    if split_label:
+        print(f"Configs selected for split {split_label}: {len(jobs)}", flush=True)
+    else:
+        print(f"Configs selected: {len(jobs)}", flush=True)
+    print(f"Node selected: {hostname}", flush=True)
+    print(f"GPUs selected: {', '.join(FIXED_GPU_TOKENS)}", flush=True)
+    print(f"Max parallel trainings: {max_parallel}", flush=True)
+    print(f"Training args: {' '.join(DEFAULT_TRAINING_ARGS + tuple(forwarded_args))}", flush=True)
+    print(f"Launcher logs: {repo_relative(log_dir)}", flush=True)
+    print("Launch mode: local queued trainings only", flush=True)
+
+    if args.test or args.dry_run:
+        print("Running launcher self-test; no trainings will be started.", flush=True)
+        for index, job in enumerate(jobs, start=1):
+            slot = execution_slots[(index - 1) % len(execution_slots)]
+            command = build_training_command(job, forwarded_args)
+            print(
+                f"[dry-run {index:02d}] slot={format_slot(slot, hostname)} "
+                f"dir={job.training_directory} cmd={shlex.join(command)}",
+                flush=True,
+            )
+        return 0
+
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    pending_jobs = deque(jobs)
+    free_slots = deque(execution_slots)
+    active_processes: list[ActiveProcess] = []
+    failed_jobs: list[tuple[TrainingJob, int, Path]] = []
+    completed_jobs = 0
+
+    try:
+        while pending_jobs or active_processes:
+            while pending_jobs and free_slots:
+                job = pending_jobs.popleft()
+                slot = free_slots.popleft()
+                active_processes.append(launch_job(job, slot, forwarded_args, log_dir, hostname))
+
+            if not active_processes:
+                break
+
+            time.sleep(args.poll_seconds)
+
+            still_active: list[ActiveProcess] = []
+            for active in active_processes:
+                return_code = active.process.poll()
+                if return_code is None:
+                    still_active.append(active)
+                    continue
+
+                active.log_handle.close()
+                free_slots.append(active.slot)
+
+                if return_code == 0:
+                    completed_jobs += 1
+                    print(
+                        f"[done] slot={format_slot(active.slot, hostname)} "
+                        f"config={repo_relative(active.job.config_path)} "
+                        f"completed={completed_jobs}/{len(jobs)}",
+                        flush=True,
+                    )
+                else:
+                    failed_jobs.append((active.job, return_code, active.log_path))
+                    print(
+                        f"[fail] slot={format_slot(active.slot, hostname)} "
+                        f"config={repo_relative(active.job.config_path)} "
+                        f"exit={return_code} log={repo_relative(active.log_path)}",
+                        flush=True,
+                    )
+                    if args.stop_on_failure:
+                        pending_jobs.clear()
+
+            active_processes = still_active
+
+            if failed_jobs and args.stop_on_failure:
+                terminate_active_processes(active_processes)
+                active_processes = []
+                break
+
+    except KeyboardInterrupt:
+        print("\nInterrupted. Terminating active trainings...", flush=True)
+        terminate_active_processes(active_processes)
+        return 130
+
+    print(f"Completed {completed_jobs} training(s); failed {len(failed_jobs)}.", flush=True)
+    if failed_jobs:
+        for job, return_code, log_path in failed_jobs:
+            print(
+                f"  {repo_relative(job.config_path)} exit={return_code} "
+                f"log={repo_relative(log_path)}",
+                flush=True,
+            )
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

root_gnn_dgl/scripts/check_dataset_files.py

@@ -2,6 +2,9 @@ import yaml
 import os
 import subprocess
 import argparse
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
 
 def check_dataset_files(yaml_file, rerun=False):
     """
@@ -69,9 +72,9 @@ def check_dataset_files(yaml_file, rerun=False):
                 # Optionally rerun data prep
                 if rerun:
                     print(f"  Reprocessing dataset '{dataset_name}' ...")
-                    prep_command = f"jobs/prep_data/prep_data.sh {yaml_file} {dataset_name} {chunks}"
+                    prep_command = [str(REPO_ROOT / "jobs/prep_data/prep_data.sh"), yaml_file, dataset_name, str(chunks)]
                     try:
-                        subprocess.run(prep_command, shell=True, check=True)
+                        subprocess.run(prep_command, check=True)
                     except subprocess.CalledProcessError as e:
                         print(f"  Could NOT reprocess '{dataset_name}': {e}")
             else:
@@ -87,8 +90,7 @@ def check_dataset_files(yaml_file, rerun=False):
         print(f"Error processing {yaml_file}: {e}")
 
 def main(pargs):
-    # Base directory containing the YAML files
-    base_directory = os.getcwd() + "/configs/"
+    base_directory = REPO_ROOT / "configs"
 
     if pargs.configs:
         configs = [p.strip() for p in pargs.configs.split(',')]
@@ -101,10 +103,11 @@ def main(pargs):
         ]
 
     for config in configs:
-        yaml_file = os.path.join(base_directory, config)
-        if os.path.exists(yaml_file):
+        config_path = Path(config).expanduser()
+        yaml_file = config_path if config_path.is_absolute() else base_directory / config_path
+        if yaml_file.exists():
             print(f"\nProcessing file: {config}")
-            check_dataset_files(yaml_file, pargs.rerun)
+            check_dataset_files(str(yaml_file), pargs.rerun)
         else:
             print(f"File not found: {yaml_file}")
 
@@ -114,7 +117,7 @@ if __name__ == "__main__":
         "--configs", "-c",
         type=str,
         required=False,
-        help="Comma-separated list of YAML config paths relative to base directory"
+        help="Comma-separated list of YAML config paths relative to configs/ or absolute paths"
     )
     parser.add_argument(
         "--rerun", "-r",
@@ -122,4 +125,4 @@ if __name__ == "__main__":
         help="Automatically re-run data processing to fix missing files"
     )
     args = parser.parse_args()
-    main(args)
+    main(args)

root_gnn_dgl/scripts/export_onnx.py

@@ -23,8 +23,8 @@ from pathlib import Path
 from types import MethodType, SimpleNamespace
 from typing import Any, Dict, Iterator, Optional, Tuple
 
-file_path = os.getcwd()
-sys.path.append(file_path)
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
 
 import dgl
 import matplotlib.pyplot as plt

root_gnn_dgl/scripts/inference.py

@@ -1,6 +1,4 @@
 import sys
-file_path = "/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl"
-sys.path.append(file_path)
 import os
 import argparse
 import yaml
@@ -8,6 +6,10 @@ import gc
 from array import array
 import fnmatch
 import re
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
 
 import torch
 import dgl

root_gnn_dgl/scripts/plot_config_distributions.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+
+import argparse
+import glob
+import math
+import os
+import re
+from pathlib import Path
+
+import awkward as ak
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_pdf import PdfPages
+import numpy as np
+import uproot
+import yaml
+
+
+OPERATORS = {
+    ">": lambda values, cut: values > cut,
+    ">=": lambda values, cut: values >= cut,
+    "<": lambda values, cut: values < cut,
+    "<=": lambda values, cut: values <= cut,
+    "==": lambda values, cut: values == cut,
+    "!=": lambda values, cut: values != cut,
+}
+
+
+def load_config(path):
+    with open(path, "r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def selection_branches(selection):
+    if isinstance(selection, str):
+        tokens = re.findall(r"\b[A-Za-z_][A-Za-z0-9_]*\b", selection)
+        keywords = {"and", "or", "not", "True", "False"}
+        return [token for token in tokens if token not in keywords]
+    if isinstance(selection, (list, tuple)) and len(selection) > 0:
+        return [selection[0]]
+    return []
+
+
+def feature_branches(node_branch_names):
+    branches = []
+    for feature in node_branch_names:
+        if not isinstance(feature, list):
+            continue
+        for branch in feature:
+            if isinstance(branch, str) and branch != "CALC_E":
+                branches.append(branch)
+    return branches
+
+
+def branches_for_dataset(dataset_config):
+    args = dataset_config["args"]
+    branches = feature_branches(args.get("node_branch_names", []))
+    for selection in dataset_config.get("selections", []):
+        branches.extend(selection_branches(selection))
+    return sorted(set(branches))
+
+
+def resolve_files(args):
+    raw_dir = args["raw_dir"]
+    file_names = args["file_names"]
+    files = []
+    if isinstance(file_names, str):
+        files.extend(glob.glob(os.path.join(raw_dir, file_names)))
+    else:
+        for file_name in file_names:
+            files.extend(glob.glob(os.path.join(raw_dir, file_name)))
+    return sorted(files)
+
+
+def load_arrays(dataset_config, max_events=None):
+    args = dataset_config["args"]
+    tree_name = args.get("tree_name", "nominal_Loose")
+    branches = branches_for_dataset(dataset_config)
+    arrays = []
+    events_read = 0
+
+    for file_name in resolve_files(args):
+        with uproot.open(file_name) as root_file:
+            tree = root_file[tree_name]
+            entry_stop = None
+            if max_events is not None:
+                remaining = max_events - events_read
+                if remaining <= 0:
+                    break
+                entry_stop = remaining
+            array = tree.arrays(branches, library="ak", entry_stop=entry_stop)
+            arrays.append(array)
+            if branches:
+                events_read += len(array[branches[0]])
+            if max_events is not None and events_read >= max_events:
+                break
+
+    if not arrays:
+        pattern = os.path.join(args["raw_dir"], str(args["file_names"]))
+        raise FileNotFoundError(f"No files found for pattern {pattern}")
+    return ak.concatenate(arrays, axis=0)
+
+
+def selection_mask(data, selections):
+    first_field = data.fields[0] if len(data.fields) > 0 else None
+    if first_field is None:
+        return None
+
+    mask = np.ones(len(data[first_field]), dtype=bool)
+    for selection in selections:
+        if isinstance(selection, str):
+            current_mask = eval(selection, {"__builtins__": {}}, data)
+        else:
+            branch, cut, op = selection
+            if op not in OPERATORS:
+                raise ValueError(f"Unknown selection operator: {op}")
+            current_mask = OPERATORS[op](data[branch], cut)
+        mask = mask & ak.to_numpy(current_mask)
+    return mask
+
+
+def ensure_node_array(value, reference):
+    if isinstance(value, (int, float, complex)):
+        return ak.full_like(reference, value)
+    return value
+
+
+def branch_array(data, branch, node_type, reference):
+    value = ensure_node_array(branch, reference)
+    if not isinstance(branch, str):
+        return value
+    value = data[branch]
+    if node_type == "single":
+        return ak.singletons(value)
+    return value
+
+
+def clean_label(value):
+    return str(value).replace("/", "_")
+
+
+def per_type_feature_label(feature_spec, type_index):
+    return clean_label(feature_spec[type_index])
+
+
+def flatten_values(values):
+    flat_values = np.asarray(ak.to_numpy(ak.ravel(values)), dtype=float)
+    return flat_values[np.isfinite(flat_values)]
+
+
+def build_feature_values(data, dataset_config):
+    args = dataset_config["args"]
+    node_branch_names = args["node_branch_names"]
+    node_branch_types = args["node_branch_types"]
+    node_feature_scales = [float(scale) for scale in args["node_feature_scales"]]
+    n_types = len(node_branch_names[0])
+
+    references = []
+    for type_index in range(n_types):
+        branch = node_branch_names[0][type_index]
+        node_type = node_branch_types[type_index]
+        if isinstance(branch, str):
+            reference = data[branch]
+            if node_type == "single":
+                reference = ak.singletons(reference)
+        else:
+            raise ValueError("The first node feature must use real branches to define node counts.")
+        references.append(reference)
+
+    features = {}
+    pt_parts = []
+    eta_parts = []
+
+    for type_index in range(n_types):
+        pt_parts.append(branch_array(data, node_branch_names[0][type_index], node_branch_types[type_index], references[type_index]))
+        eta_parts.append(branch_array(data, node_branch_names[1][type_index], node_branch_types[type_index], references[type_index]))
+
+    for feature_index, feature_spec in enumerate(node_branch_names):
+        if not isinstance(feature_spec, list):
+            continue
+
+        per_type_parts = []
+        for type_index in range(n_types):
+            branch = feature_spec[type_index]
+            if not isinstance(branch, str) or branch == "CALC_E":
+                per_type_parts.append(None)
+                continue
+
+            reference = references[type_index]
+            node_type = node_branch_types[type_index]
+            per_type_parts.append(branch_array(data, branch, node_type, reference))
+
+        for type_index, part in enumerate(per_type_parts):
+            if part is None:
+                continue
+            scaled_part = part * node_feature_scales[feature_index]
+            scaled_pt = pt_parts[type_index] * node_feature_scales[0]
+            scaled_part = scaled_part[scaled_pt != 0]
+            features[per_type_feature_label(feature_spec, type_index)] = flatten_values(scaled_part)
+
+    return features
+
+
+def common_bins(datasets, feature_name, bins):
+    values = np.concatenate([features[feature_name] for features in datasets.values() if len(features[feature_name]) > 0])
+    if len(values) == 0:
+        return np.linspace(0, 1, bins + 1)
+
+    unique_values = np.unique(values)
+    if len(unique_values) <= 20 and np.allclose(unique_values, np.round(unique_values)):
+        low = math.floor(values.min())
+        high = math.ceil(values.max())
+        return np.arange(low - 0.5, high + 1.5, 1)
+
+    low, high = np.percentile(values, [0.5, 99.5])
+    if not np.isfinite(low) or not np.isfinite(high) or low == high:
+        low, high = values.min(), values.max()
+    if low == high:
+        low -= 0.5
+        high += 0.5
+    return np.linspace(low, high, bins + 1)
+
+
+def plot_feature(feature_name, datasets, bins):
+    fig, ax = plt.subplots(figsize=(8, 6))
+    hist_bins = common_bins(datasets, feature_name, bins)
+
+    for dataset_name, features in datasets.items():
+        values = features[feature_name]
+        if len(values) == 0:
+            continue
+        ax.hist(
+            values,
+            bins=hist_bins,
+            histtype="step",
+            density=True,
+            linewidth=1.8,
+            label=f"{dataset_name} (n={len(values)})",
+        )
+
+    ax.set_xlabel(feature_name)
+    ax.set_ylabel("Normalized entries")
+    ax.legend(frameon=False)
+    ax.grid(alpha=0.25)
+    fig.tight_layout()
+    return fig
+
+
+def safe_filename(name):
+    return re.sub(r"[^A-Za-z0-9_.-]+", "_", name).strip("_")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Plot model-input node feature distributions for every dataset in a config."
+    )
+    parser.add_argument("--config", required=True, help="YAML config containing Datasets.")
+    parser.add_argument(
+        "--output-dir",
+        default=None,
+        help="Directory for optional PNG outputs. Defaults to plots/<config-stem>_distributions.",
+    )
+    parser.add_argument(
+        "--output-pdf",
+        default=None,
+        help="Path for the multi-page PDF. Defaults to plots/<config-stem>_distributions.pdf.",
+    )
+    parser.add_argument("--write-pngs", action="store_true", help="Also write one PNG per plot.")
+    parser.add_argument("--bins", type=int, default=80, help="Number of bins for continuous features.")
+    parser.add_argument("--max-events", type=int, default=None, help="Optional maximum events per dataset.")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    output_dir = Path(args.output_dir) if args.output_dir else Path("plots") / f"{Path(args.config).stem}_distributions"
+    output_pdf = Path(args.output_pdf) if args.output_pdf else Path("plots") / f"{Path(args.config).stem}_distributions.pdf"
+    output_pdf.parent.mkdir(parents=True, exist_ok=True)
+    if args.write_pngs:
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    dataset_features = {}
+    for dataset_name, dataset_config in config["Datasets"].items():
+        print(f"Loading {dataset_name}", flush=True)
+        data = load_arrays(dataset_config, max_events=args.max_events)
+        mask = selection_mask(data, dataset_config.get("selections", []))
+        if mask is not None:
+            data = data[mask]
+        dataset_features[dataset_name] = build_feature_values(data, dataset_config)
+
+    feature_names = list(next(iter(dataset_features.values())).keys())
+    with PdfPages(output_pdf) as pdf:
+        for feature_name in feature_names:
+            fig = plot_feature(feature_name, dataset_features, args.bins)
+            pdf.savefig(fig)
+            if args.write_pngs:
+                output_path = output_dir / f"{safe_filename(feature_name)}.png"
+                fig.savefig(output_path, dpi=160)
+                print(f"Wrote {output_path}", flush=True)
+            plt.close(fig)
+    print(f"Wrote {output_pdf}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

root_gnn_dgl/scripts/prep_data.py

@@ -1,8 +1,10 @@
 import sys
 import os
 import glob
-file_path = os.getcwd()
-sys.path.append(file_path)
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
 
 import root_gnn_base.utils as utils
 from root_gnn_base.dataset import compute_cutflow, print_cutflow

root_gnn_dgl/scripts/selections.py

@@ -2,12 +2,13 @@ import argparse
 import glob
 import os
 import sys
+from pathlib import Path
 
 import awkward as ak
 import uproot
 
-file_path = os.getcwd()
-sys.path.append(file_path)
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
 
 from root_gnn_base import utils
 from root_gnn_base.dataset import selection_branches, check_selection, print_cutflow, init_cutflow

root_gnn_dgl/scripts/training_script.py

@@ -3,6 +3,7 @@ import time
 import datetime
 import yaml
 import os
+from pathlib import Path
 
 start_time = time.time()
 
@@ -11,8 +12,8 @@ import torch
 import torch.nn as nn
 
 import sys
-file_path = os.getcwd()
-sys.path.append(file_path)
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
 import root_gnn_base.batched_dataset as datasets
 from root_gnn_base import utils
 import root_gnn_base.custom_scheduler as lr_utils