eqcctpro 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

eqcctpro/tools.py

@@ -0,0 +1,968 @@
+"""
+tools.py contains the sub-tool functions that EQCCTPro uses, such as getting the current system's VRAM, creating subdirs, etc. 
+"""
+from __future__ import annotations
+import os 
+import re
+import ast
+import math
+import glob
+import csv
+import shutil 
+import pynvml 
+import logging
+import pandas as pd 
+from math import floor
+from pathlib import Path
+from dataclasses import dataclass, field # To-Do: Figure out what dataclass and field mean/do
+from datetime import datetime, timedelta 
+from typing import Iterable, Callable, Optional, List
+
+# We define the CSV header as a globally consistent CSV header for all CSVs generated by EQCCTPro
+CANONICAL_CSV_HEADER = [
+    "Trial Number",
+    "Stations Used",
+    "Number of Stations Used",
+    "Number of CPUs Allocated for Ray to Use",
+    "Intra-parallelism Threads",
+    "Inter-parallelism Threads",
+    "GPUs Used",
+    "Inference Actor Memory Limit (MB)",
+    "Total Waveform Analysis Timespace (min)",
+    "Total Number of Timechunks",
+    "Concurrent Timechunks Used",
+    "Length of Timechunk (min)",
+    "Number of Concurrent Station Tasks",
+    "Total Run time for Picker (s)",
+    "Model Used",
+    "Trial Success",
+    "Error Message",
+]
+
+"""
+build_station_list_from_dir, looks_like_timechunk_id, and _TIMECHUNK_RE work together to discover the stations that are in a timechunk directory. \
+build_station_list_from_dir builds the list of stations in a timechunk dir, using looks_like_timechunk_id to make sure that the subdirectory of the 
+given input_dir is not a timechunk_id, if it is, it is filtered out. 
+"""
+# Pattern for understanding the structure of a timechunk directory naming pattern 
+_TIMECHUNK_RE = re.compile(r"^\d{8}T\d{6}Z_\d{8}T\d{6}Z$")
+
+def looks_like_timechunk_id(name: str) -> bool:
+    return bool(_TIMECHUNK_RE.match(name or ""))
+
+def build_station_list_from_dir(input_dir: str) -> list[str]:
+    """
+    Robustly discover stations under a timechunk directory.
+    Accepts files like *.mseed/*.sac or one-dir-per-station structures.
+    """
+    stations = set()
+
+    # 1) Files directly inside input_dir
+    for p in glob.glob(os.path.join(input_dir, "*")):
+        base = os.path.basename(p)
+        if os.path.isfile(p):
+            # file path — take stem without extension
+            stations.add(os.path.splitext(base)[0])
+
+    # 2) One subdir per station (e.g., input_dir/AT01/*.mseed)
+    for p in glob.glob(os.path.join(input_dir, "*")):
+        if os.path.isdir(p):
+            stations.add(os.path.basename(p))
+
+    # Filter out anything that looks like a timechunk id (safety)
+    stations = [s for s in stations if not looks_like_timechunk_id(s)]
+
+    return sorted(stations)
+
+"""
+generate_station_list builds a list that contains the number of stations to use in the trial iterations, marching from a given amount of starting stations 
+to a final amount of stations with either an specified or a pre-determined step size.
+"""
+def generate_station_list(starting_amount_of_stations, total_num_stations_to_use, station_list_step_size):
+    if total_num_stations_to_use == 1:
+        return [1]
+    elif total_num_stations_to_use < 10:
+        return list(range(1, total_num_stations_to_use + 1))
+    elif total_num_stations_to_use >= 10 and starting_amount_of_stations == 1 and station_list_step_size == 1:
+        # We want to reduce as much extra testing as needed by applying an iterative approach for marching approach via smart step sizes
+        
+        # We check if we can apply multiples of 5 up to total_num_stations_to_use
+        target = total_num_stations_to_use
+        start = 10 # We start from 10 because we already covered 1-9 in the previous condition
+        step = 5 # Step size of 5 for multiples of 5 (User can change this if desired - To Do add var for this in the future)
+        
+        stp_of_one = list(range(1, start + 1)) # We generate a list of 1-10 with step size of 1 and afterwards we will add multiples of 5
+
+        remainder = target % step # We calculate the remainder to see if target is a multiple of 5 or not
+        max_multiple = target - remainder # We calculate what is the maximum multiple of 5 that is less than or equal to the target
+        if remainder != 0:  # We found a number that is not immediately a multiple of 5, so we need to add the remaining numbers after the last multiple of 5 up to the target with step size of 1 
+            marching_scheme = stp_of_one + list(range(start + step, max_multiple + 1, step)) + list(range(max_multiple + 1, target + 1, 1))
+
+        elif remainder == 0: # We hit that the target is not a multiple of 5, so we start from 20, ..., +5, ..., up to target
+            marching_scheme = stp_of_one + list(range(start + step, target + 1, step))
+
+        return sorted(set(marching_scheme))
+    else: 
+        return list(range(starting_amount_of_stations, total_num_stations_to_use + 1, station_list_step_size))
+
+"""
+list_gpu_ids returns a list of available GPU IDs on the system.
+"""
+def list_gpu_ids():
+    """List all available GPU IDs on the system."""
+    pynvml.nvmlInit()  # Initialize NVML
+    gpu_count = pynvml.nvmlDeviceGetCount()  # Get number of GPUs
+    gpu_ids = list(range(gpu_count))  # Create a list of GPU indices
+    pynvml.nvmlShutdown()  # Shutdown NVML
+    return gpu_ids  
+
+"""
+get_gpu_vram retrieves the total and free VRAM (in GB) for the current GPU.
+"""
+def get_gpu_vram(gpu_index): # Need to fix to pass in a GPU instead of hard indexing 0 
+    """Retrieve total and free VRAM (in GB) for the current GPU."""
+    pynvml.nvmlInit()  # Initialize NVML
+    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_index)  # Use provided GPU
+    total_vram = pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024**3)  # Convert bytes to GB
+    free_vram = pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024**3)  # Convert bytes to GB
+    pynvml.nvmlShutdown()  # Shutdown NVML
+    return total_vram, free_vram        
+
+"""
+_parse_gpu_fields normalize the 'GPUs Used' field from CSV into a list[int]
+"""
+def _parse_gpus_field(x):
+    if x is None:
+        return []
+    if isinstance(x, (list, tuple)):
+        return [int(v) for v in x]
+    if isinstance(x, (int, float)) and not (isinstance(x, float) and math.isnan(x)):
+        return [int(x)]
+    if isinstance(x, str):
+        s = x.strip()
+        if not s:
+            return []
+        s = s.replace("(", "[").replace(")", "]").replace("{", "[").replace("}", "]")
+        try:
+            val = ast.literal_eval(s)
+            if isinstance(val, (list, tuple)):
+                return [int(v) for v in val]
+            if isinstance(val, (int, float)):
+                return [int(val)]
+        except Exception:
+            nums = re.findall(r"-?\d+", s)
+            return [int(n) for n in nums]
+    return []
+
+"""
+VramPlan checks to see that for the RunEQCCTPro functionality, when the user calls to use a certain amount of VRAM per Raylet,
+we check to see that the amount requested does not exceed what is capable for the system. 
+"""
+# To-Do: Understand how this works
+@dataclass
+class VramPlan:
+    intended_workers: int
+    per_worker_mb: float
+    overhead_mb: float
+    model_vram_mb: float
+    per_gpu_cap_mb: List[float]         # post safety-cap
+    max_workers_per_gpu: List[int]
+    total_max_workers: int
+    aggregate_cap_mb: float
+    aggregate_need_mb: float
+    ok_per_gpu: bool
+    ok_aggregate: bool
+    # Optional; value-object semantics preserved (not shown in repr/eq)
+    logger: Optional[logging.Logger] = field(default=None, repr=False, compare=False)
+
+def evaluate_vram_capacity(
+    *,
+    intended_workers: int,
+    vram_per_worker_mb: float,
+    per_gpu_free_mb: Iterable[float],        # free MB per selected GPU (raw)
+    model_vram_mb: float = 3000.0,           # reserve per GPU
+    safety_cap: float = 0.90,                # cap each GPU
+    eqcct_overhead_gb: float = 1.1           # per-worker overhead (GB)
+) -> VramPlan:
+    """
+    Computes both per-GPU admission and aggregate budget feasibility.
+    Pure: no logging/printing/raising.
+    """
+    overhead_mb = eqcct_overhead_gb * 1024.0
+    per_gpu_cap_mb = [safety_cap * float(x) for x in per_gpu_free_mb]
+
+    denom = (vram_per_worker_mb + overhead_mb)
+    if denom <= 0:
+        raise ValueError("Non-positive per-worker memory requirement.")
+
+    max_workers_per_gpu = []
+    for cap in per_gpu_cap_mb:
+        remaining = max(0.0, cap - model_vram_mb)  # reserve the model actor once per GPU
+        maxw = max(0, floor(remaining / denom))
+        max_workers_per_gpu.append(maxw)
+
+    total_max_workers = sum(max_workers_per_gpu)
+    ok_per_gpu = intended_workers <= total_max_workers
+
+    aggregate_cap_mb = sum(per_gpu_cap_mb)
+    aggregate_need_mb = (model_vram_mb * len(per_gpu_cap_mb)) + intended_workers * denom
+    ok_aggregate = aggregate_need_mb <= aggregate_cap_mb
+
+    return VramPlan(
+        intended_workers=intended_workers,
+        per_worker_mb=vram_per_worker_mb,
+        overhead_mb=overhead_mb,
+        model_vram_mb=model_vram_mb,
+        per_gpu_cap_mb=per_gpu_cap_mb,
+        max_workers_per_gpu=max_workers_per_gpu,
+        total_max_workers=total_max_workers,
+        aggregate_cap_mb=aggregate_cap_mb,
+        aggregate_need_mb=aggregate_need_mb,
+        ok_per_gpu=ok_per_gpu,
+        ok_aggregate=ok_aggregate,
+    )
+
+def _emit(logger: Optional[logging.Logger], level: str, msg: str) -> None:
+    if logger is not None:
+        getattr(logger, level)(msg)
+    else:
+        print(msg)
+
+# ---------- Thin wrappers to preserve your two call sites ----------
+
+def check_vram_per_gpu_style(
+    *,
+    selected_gpus: List[int],
+    get_gpu_vram_fn: Callable[[int], tuple[float, float]],  # returns (total_gb, free_gb) for that GPU
+    intended_workers: int,
+    vram_mb: float,
+    model_vram_mb: float = 3000.0,
+    safety_cap: float = 0.95,       # matches your original 95% cap
+    eqcct_overhead_gb: float = 0.0, # original per-GPU check ignored runtime overhead
+    logger: Optional[logging.Logger] = None
+) -> None:
+    
+    safety_cap = float(safety_cap)
+    if not (0.0 < safety_cap <= 0.99):
+        raise ValueError(f"safety_cap must be in (0, 0.99], got {safety_cap}.")
+    
+    per_gpu_free_mb = [(get_gpu_vram_fn(gid)[1] * 1024.0) for gid in selected_gpus]
+
+    plan = evaluate_vram_capacity(
+        intended_workers=intended_workers,
+        vram_per_worker_mb=float(vram_mb),
+        per_gpu_free_mb=per_gpu_free_mb,
+        model_vram_mb=model_vram_mb,
+        safety_cap=safety_cap,
+        eqcct_overhead_gb=eqcct_overhead_gb,
+    )
+    plan.logger = logger
+
+    if not plan.ok_per_gpu:
+        # Compose precise diagnostic
+        unit = plan.per_worker_mb + plan.overhead_mb
+        msg = (
+            f"ERROR: Per-GPU capacity insufficient for {plan.intended_workers} workers.\n"
+            f"  Reservation per GPU: model={plan.model_vram_mb:.0f} MB, worker_unit={unit:.0f} MB "
+            f"({plan.per_worker_mb:.0f} + {plan.overhead_mb:.0f})\n"
+            f"  Per-GPU caps after safety: {', '.join(f'{c:.0f}' for c in plan.per_gpu_cap_mb)} MB\n"
+            f"  Max workers per GPU: {plan.max_workers_per_gpu} (total={plan.total_max_workers})\n"
+            f"Action: lower vram_mb, reduce concurrency, or add GPUs."
+        )
+        _emit(logger, "error", msg)
+        raise RuntimeError(msg)
+
+    _emit(logger, "info",
+          f"Per-GPU admission OK: {plan.intended_workers} ≤ {plan.total_max_workers} workers.")
+
+def check_vram_aggregate_style(
+    *,
+    eval_mode: str,
+    selected_gpus: List[int],
+    get_cluster_free_gb_fn: Callable[[], tuple[float, float]], # returns (total_gb, free_gb)
+    intended_workers: int,
+    vram_mb: float,
+    model_vram_mb: float = 3000.0,
+    safety_cap: float = 0.90,
+    eqcct_overhead_gb: float = 1.1,
+    logger: Optional[logging.Logger] = None
+) -> None:
+    if eval_mode.lower() != "gpu":
+        raise ValueError(f"vram_mb is only meaningful in GPU mode; got eval_mode='{eval_mode}'.")
+
+    # Homogeneous assumption (your prior logic): same free_gb replicated
+    _, free_gb = get_cluster_free_gb_fn()
+    per_gpu_free_mb = [free_gb * 1024.0] * len(selected_gpus)
+
+    plan = evaluate_vram_capacity(
+        intended_workers=intended_workers,
+        vram_per_worker_mb=float(vram_mb),
+        per_gpu_free_mb=per_gpu_free_mb,
+        model_vram_mb=model_vram_mb,
+        safety_cap=safety_cap,
+        eqcct_overhead_gb=eqcct_overhead_gb,
+    )
+    plan.logger = logger
+
+    if not plan.ok_aggregate:
+        unit = plan.per_worker_mb + plan.overhead_mb
+        msg = (
+            f"ERROR: Aggregate VRAM insufficient.\n"
+            f"  GPUs: {len(selected_gpus)} | Safety cap: {int(safety_cap*100)}%\n"
+            f"  Aggregate cap: {plan.aggregate_cap_mb:.0f} MB\n"
+            f"  Aggregate need: {plan.aggregate_need_mb:.0f} MB "
+            f"(= {model_vram_mb:.0f}×{len(selected_gpus)} + {intended_workers}×{unit:.0f})\n"
+            f"Action: lower vram_mb, reduce concurrency, or add GPUs."
+        )
+        _emit(logger, "error", msg)
+        raise RuntimeError(msg)
+
+    _emit(logger, "info",
+          f"Aggregate budget OK. Need {plan.aggregate_need_mb:.0f} MB ≤ "
+          f"Cap {plan.aggregate_cap_mb:.0f} MB across {len(selected_gpus)} GPU(s).")
+
+
+"""
+prepare_csv either loads or initializes the CSV file for storing test results.
+"""      
+def prepare_csv(csv_file_path, logger):
+    """
+    Loads or initializes the CSV file for storing test results.
+    """
+    if os.path.exists(csv_file_path):
+        logger.info(f"Loading existing CSV file from '{csv_file_path}'...")
+        return pd.read_csv(csv_file_path)
+    logger.info(f"CSV file not found. Creating a new CSV file at '{csv_file_path}'...")
+    
+    columns = CANONICAL_CSV_HEADER
+    df = pd.DataFrame(columns=columns)
+    df.to_csv(csv_file_path, index=False)
+
+"""
+append_trial_row appends a completed trial_dictionary to a CSV via the last line. Does this after successfully completing a trial during the 
+EvaluateSystem process for either CPU or GPU. Does not provide information for the success or the error message of the trial (update_csv does).
+"""
+def append_trial_row(csv_path: str, trial_data: dict):
+    """
+    Append a complete trial row to the CSV with all fields populated.
+    Ensures "GPUs Used" column is consistently formatted and quoted.
+    """
+    csvp = Path(csv_path) 
+    
+    # Ensure header exists with canonical order
+    if not csvp.exists():
+        pd.DataFrame(columns=CANONICAL_CSV_HEADER).to_csv(csvp, index=False)
+
+    df_existing = pd.read_csv(csvp, keep_default_na=False)
+    
+    # Align row to the canonical header (use empty string for missing keys)
+    row = {col: trial_data.get(col, "") for col in CANONICAL_CSV_HEADER}
+    
+    # Normalize "GPUs Used" format to ensure consistent quoting
+    if "GPUs Used" in row and row["GPUs Used"]:
+        gpus_value = row["GPUs Used"]
+        # If it's already a JSON string, parse and reformat consistently
+        if isinstance(gpus_value, str):
+            try:
+                # Try to parse as JSON/list
+                gpus_list = ast.literal_eval(gpus_value.replace("(", "[").replace(")", "]"))
+                if isinstance(gpus_list, (list, tuple)):
+                    # Format consistently: [0] or [0, 1] with space after comma
+                    row["GPUs Used"] = "[" + ", ".join(map(str, gpus_list)) + "]"
+                else:
+                    row["GPUs Used"] = str(gpus_value)
+            except:
+                row["GPUs Used"] = str(gpus_value)
+        elif isinstance(gpus_value, (list, tuple)):
+            row["GPUs Used"] = "[" + ", ".join(map(str, gpus_value)) + "]"
+        else:
+            row["GPUs Used"] = str(gpus_value) if gpus_value else "[]"
+    
+    # Auto-number trials if not provided
+    if pd.isna(row["Trial Number"]) or row["Trial Number"] == "" or row["Trial Number"] is None:
+        row["Trial Number"] = len(df_existing) + 1
+
+    df_new = pd.DataFrame([row], columns=CANONICAL_CSV_HEADER)
+    df_out = pd.concat([df_existing, df_new], ignore_index=True)
+    
+    # Write CSV with quoting to ensure "GPUs Used" is always quoted
+    # Use QUOTE_NONNUMERIC to quote all non-numeric fields, ensuring consistency
+    df_out.to_csv(csvp, index=False, quoting=csv.QUOTE_NONNUMERIC)
+    
+    print(f"Appended trial {row['Trial Number']} to {csv_path}")
+
+"""
+normalize_gpu_csv_quoting ensures all "GPUs Used" entries in a CSV are consistently formatted and quoted.
+This function can be used to update existing CSV files to have consistent formatting.
+"""
+def normalize_gpu_csv_quoting(csv_filepath: str):
+    """
+    Normalize the "GPUs Used" column in an existing CSV file to ensure consistent formatting and quoting.
+    """
+    if not os.path.exists(csv_filepath):
+        return
+    
+    df = pd.read_csv(csv_filepath, keep_default_na=False)
+    
+    if "GPUs Used" not in df.columns:
+        return
+    
+    # Normalize each "GPUs Used" entry
+    for idx in df.index:
+        gpus_value = df.at[idx, "GPUs Used"]
+        if pd.isna(gpus_value) or gpus_value == "":
+            df.at[idx, "GPUs Used"] = "[]"
+        else:
+            try:
+                # Parse the GPU value (handles various formats)
+                gpus_str = str(gpus_value).strip()
+                # Remove existing quotes if present
+                if gpus_str.startswith('"') and gpus_str.endswith('"'):
+                    gpus_str = gpus_str[1:-1]
+                
+                # Parse as list
+                gpus_list = ast.literal_eval(gpus_str.replace("(", "[").replace(")", "]"))
+                if isinstance(gpus_list, (list, tuple)):
+                    # Format consistently: [0] or [0, 1] with space after comma
+                    df.at[idx, "GPUs Used"] = "[" + ", ".join(map(str, gpus_list)) + "]"
+                else:
+                    df.at[idx, "GPUs Used"] = str(gpus_value)
+            except:
+                # If parsing fails, keep original but ensure it's a string
+                df.at[idx, "GPUs Used"] = str(gpus_value)
+    
+    # Write back with consistent quoting
+    df.to_csv(csv_filepath, index=False, quoting=csv.QUOTE_NONNUMERIC)
+
+"""
+update_csv updates a completed trial after the code has exited the mseed_predictor loop and into the last steps of completing the trial. 
+If the trial either was a success or had errors, we update the last row with the success/error information of the trial. 
+"""
+def update_csv(csv_filepath, success, error_message):
+    df = pd.read_csv(csv_filepath)
+    if "Error Message" not in df.columns:
+        df["Error Message"] = ""
+
+    # Ensure string dtype
+    df["Error Message"] = df["Error Message"].astype("string")
+
+    last_idx = df.index[-1] # Get last row id number
+    df.loc[last_idx, 'Trial Success'] = success # Access value at row last_idx, column 'Trial Success' 
+    df.loc[last_idx, 'Error Message'] = error_message # Access value at row last_idx, column 'Error Message'
+
+    df.to_csv(csv_filepath, index=False)
+
+# """
+# remove_directory removes a specified directory if it exists. (NOT CURRENTLY USED)
+# """
+# def remove_directory(path):
+#     """
+#     Removes the specified directory if it exists.
+#     """
+#     if os.path.exists(path):
+#         shutil.rmtree(path)
+#         print(f"Removed directory: {path}")
+#     else:
+#         print(f"Directory '{path}' does not exist anymore.")
+        
+"""
+remove_output_subdirs removes all the subdirectoreies of a specified directory. We use it in the code for removing the contents of the 'output' dir
+once a trial has completed, because if we do not, then the code will recognize the directory name and believe that the waveform was already analyzed
+when in reality it was in the previous configuration iteration. Does not remove trial contents or log, just station subdirs created by mseed_predictor
+and parallel_predict.
+"""
+def remove_output_subdirs(output_dir: str, logger: logging.Logger | None = None) -> None:
+    """
+    Delete any *_outputs subdirectories in `output_dir`. 
+    Logs via `logger` if provided; otherwise falls back to print.
+    """
+    try:
+        for name in os.listdir(output_dir):
+            path = os.path.join(output_dir, name)
+            if os.path.isdir(path) and name.endswith("_outputs"):
+                shutil.rmtree(path, ignore_errors=True)
+                msg = f"Removed subdirectory: {path}"
+                (logger.info if logger else print)(msg)
+    except Exception as e:
+        msg = f"Failed to remove output subdirs in {output_dir}: {e}"
+        (logger.error if logger else print)(msg)
+
+"""
+We need to check to make sure that the input dir's contents (IE. the station subdirs in each timechunk dir) 
+are the same length. We need to ensure that they are
+"""
+def check_station_dirs(input_dir):
+    subdir_lens, station_list_f = [], []
+    sorted_input_dir = sorted(os.listdir(input_dir))
+    subdirs = [item for item in sorted_input_dir if os.path.isdir(os.path.join(input_dir, item))] # Inherits the sorted
+    for timechunk_dir in subdirs: 
+        subdir_path = os.path.join(input_dir, timechunk_dir)
+        try: 
+            station_list = os.listdir(subdir_path)
+            subdir_lens.append(len(station_list))
+            station_list_f = station_list
+        except OSError as e:
+            print(f"Warning: Could not read directory {subdir_path}. Error: {e}")
+
+    check_if_lens_are_same = len(set(subdir_lens))   
+    if check_if_lens_are_same != 1:
+        # Bad case
+        statement = f"The contents across your timechunk directories are not the same. They must match for EQCCTPro. Fix station subdirs so each timechunk has the same stations. Exiting..."
+        return statement, station_list_f, True
+    else: 
+        statement = f"Stations subdirs in timechunk directories are consistent. Continuing EQCCTPro..."     
+        return statement, station_list_f, False 
+
+
+"""
+tf_environ sets the tensorflow environment for either an allocated CPU or GPU configuration. 
+If it is a CPU, the CUDA_DEVICE_ORDER is configured for a CPU; Intra/inter parallelism threads are also set as well
+If a GPU, the CUDA_DEVICE_ORDER is set for a GPU, as well as intra/inter threads, and we configure that GPU to only 
+use up to a limited number of vram (vram_limit_mb), which we need for the trials.
+"""
+def tf_environ(gpu_id, vram_limit_mb=None, gpus_to_use=None, intra_threads=None, inter_threads=None, log_device=True, logger=None, skip_tf=False):
+    """
+    Configure TensorFlow to use fixed VRAM slices per visible GPU.
+    Call this ONCE per Ray actor, BEFORE building/loading any TF model.
+    """
+
+    # Normalize logger: if None, use a silent logger that discards records
+    if logger is None:
+        logger = logging.getLogger("eqcctpro.null")
+        logger.propagate = False
+        if not logger.handlers:
+            # logger.addHandler(logging.StreamHandler())
+            logger.addHandler(logging.NullHandler())
+
+    # C++ backend verbosity (must be set before importing TF)
+    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")  # 0=all,1=INFO-,2=WARNING-,3=ERROR-
+    os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") # avoid oneDNN “custom ops” info line
+
+    # 0) Visibility must be set BEFORE importing tensorflow
+    if gpu_id == -1:
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+        logger.info(f"GPU disabled (CPU-only).")
+    elif gpus_to_use is not None:
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus_to_use))
+        logger.info(f"GPU enabled. Visible GPU IDs: {gpus_to_use}")
+    else:
+        logger.info(f"GPU visibility left to environment (Ray). CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}")
+
+    if skip_tf:
+        logger.info("Skipping framework-specific initialization in the driver process.")
+        return
+
+    # 1) Now import TF (it will honor visibility)
+    import tensorflow as tf
+    tf.get_logger().setLevel(logging.ERROR)
+    try:
+        from absl import logging as absl_logging
+        absl_logging.set_verbosity(absl_logging.ERROR)
+    except Exception:
+        pass
+    if log_device:
+        tf.debugging.set_log_device_placement(True)
+
+    # 2) Threading (optional)
+    if intra_threads is not None:
+        tf.config.threading.set_intra_op_parallelism_threads(int(intra_threads))
+        logger.info(f"Configured Intra-op threads = {intra_threads}")
+    if inter_threads is not None:
+        tf.config.threading.set_inter_op_parallelism_threads(int(inter_threads))
+        logger.info(f"Configured Inter-op threads = {inter_threads}")
+
+    # 3) Configure fixed VRAM slices on all visible GPUs
+    vis_gpus = tf.config.list_physical_devices("GPU")
+    if not vis_gpus:
+        logger.info(f"TensorFlow: No GPUs visible; TF will proceed on CPU.")
+        logger.info("")
+        return {"logical_gpus": [], "physical_gpus": []}
+
+    if vram_limit_mb is None or vram_limit_mb <= 0:
+        raise ValueError("vram_limit_mb must be a positive integer when using fixed VRAM slicing.")
+
+    try:
+        for gpu in vis_gpus:
+            # One logical device per physical GPU, each with a hard VRAM cap
+            tf.config.set_logical_device_configuration(gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=int(vram_limit_mb))])
+        # Force logical devices to materialize
+        # Logical devices are a virtual rep. of a physical hardware component (CPU/GPU) that TF creates to manage workload distribution
+        # Can have more logical devices than what you have physically, however you are constrained by the physical limitations of your hardware 
+        logical = tf.config.list_logical_devices("GPU") 
+        logger.info(
+            f"Set VRAM slicing: "
+            f"{vram_limit_mb} MB per logical GPU "
+            f"({len(logical)} logical over {len(vis_gpus)} physical)."
+        )
+    except RuntimeError as e:
+        # Happens if any TF GPU context was already initialized
+        raise RuntimeError(
+            "Failed to set logical device configuration. "
+            "Ensure tf_environ() is called before any TensorFlow GPU ops or model creation.\n"
+            f"Original error: {e}")
+
+"""
+find_optimal_configurations_cpu/gpu find at the end of the trial run throught find the: 
+    1. The best number of concurrent predictions for each (stations, CPUs/GPU VRAM) pair that results in the fastest runtime.
+    2. The overall best configuration balancing stations, CPUs/GPU VRAM, and runtime.
+"""
+def find_optimal_configurations_cpu(df):
+    """
+    Find:
+    1. The best number of concurrent predictions for each (stations, CPUs) pair that results in the fastest runtime.
+    2. The overall best configuration balancing stations, CPUs, and runtime.
+    """
+
+    # Convert relevant columns to numeric, handling NaNs gracefully
+    df["Number of Stations Used"] = pd.to_numeric(df["Number of Stations Used"], errors="coerce")
+    df["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df["Number of CPUs Allocated for Ray to Use"], errors="coerce")
+    df["Total Number of Timechunks"] = pd.to_numeric(df["Total Number of Timechunks"], errors="coerce")
+    df["Concurrent Timechunks Used"] = pd.to_numeric(df["Concurrent Timechunks Used"], errors="coerce")
+    df["Number of Concurrent Station Tasks"] = pd.to_numeric(df["Number of Concurrent Station Tasks"], errors="coerce")
+    df["Total Run time for Picker (s)"] = pd.to_numeric(df["Total Run time for Picker (s)"], errors="coerce")
+    
+
+    # Drop rows with missing values in these essential columns
+    df_cleaned = df.dropna(subset=["Number of Stations Used", "Number of CPUs Allocated for Ray to Use", 
+                                "Concurrent Timechunks Used", "Number of Concurrent Station Tasks", "Total Run time for Picker (s)"])
+
+    # Find the best concurrent prediction configuration for each combination of (Stations, Timechunks, CPUs, Model)
+    optimal_concurrent_preds = df_cleaned.loc[
+        df_cleaned.groupby(["Number of Stations Used", "Concurrent Timechunks Used", "Number of CPUs Allocated for Ray to Use", "Model Used"])
+        ["Total Run time for Picker (s)"].idxmin()
+    ]
+
+    # Define what "moderate" means in terms of CPU usage (e.g., middle 50% of available CPUs)
+    cpu_min = df_cleaned["Number of CPUs Allocated for Ray to Use"].quantile(0.25)
+    cpu_max = df_cleaned["Number of CPUs Allocated for Ray to Use"].quantile(0.75)
+
+    # Filter for rows within the moderate CPU range
+    df_moderate_cpus = df_cleaned[(df_cleaned["Number of CPUs Allocated for Ray to Use"] >= cpu_min) & 
+                                (df_cleaned["Number of CPUs Allocated for Ray to Use"] <= cpu_max)]
+
+    # Sort by the highest number of stations first, then by the fastest runtime
+    best_overall_config = df_moderate_cpus.sort_values(
+        by=["Number of Stations Used", "Total Run time for Picker (s)"], 
+        ascending=[False, True]  # Maximize stations, minimize runtime
+    ).iloc[0]
+    
+    # Format the output for human readability
+    formatted_output = {
+        "Trial Number": best_overall_config["Trial Number"],
+        "Number of Stations Used": best_overall_config["Number of Stations Used"],
+        "Total Number of Timechunks": best_overall_config["Total Number of Timechunks"],
+        "Concurrent Timechunks Used": best_overall_config["Concurrent Timechunks Used"],
+        "Length of Timechunk (min)": str(best_overall_config["Length of Timechunk (min)"]),
+        "Total Waveform Analysis Timespace (min)": str(best_overall_config["Total Waveform Analysis Timespace (min)"]),
+        "Number of Concurrent Station Tasks per Timechunk": best_overall_config["Number of Concurrent Station Tasks"],
+        "Number of CPUs Allocated for Ray to Use": best_overall_config["Number of CPUs Allocated for Ray to Use"],
+        "Intra-parallelism Threads": best_overall_config["Intra-parallelism Threads"],
+        "Inter-parallelism Threads": best_overall_config["Inter-parallelism Threads"],
+        "Total Run time for Picker (s)": best_overall_config["Total Run time for Picker (s)"],
+        "Model Used": best_overall_config["Model Used"],
+        "Trial Success": best_overall_config["Trial Success"],
+        "Error Message": best_overall_config["Error Message"],
+    }
+    
+    best_overall_df = pd.DataFrame([formatted_output])
+
+
+    return optimal_concurrent_preds, best_overall_df
+
+
+def find_optimal_configurations_gpu(df):
+    """
+    Find:
+      1) Best concurrency for each (stations, CPUs, GPUs, VRAM) combo (fastest runtime).
+      2) Best overall balanced configuration.
+    """
+    # 1) Numeric normalization
+    numeric_cols = [
+        "Number of Stations Used",
+        "Number of CPUs Allocated for Ray to Use",
+        "Number of Concurrent Station Tasks",
+        "Total Run time for Picker (s)",
+        "Inference Actor Memory Limit (MB)",
+    ]
+    for col in numeric_cols:
+        df[col] = pd.to_numeric(df[col], errors="coerce")
+
+    # 2) Normalize GPUs Used -> list[int], then create a *hashable* key
+    df["GPUs Used"] = df["GPUs Used"].apply(_parse_gpus_field)
+    df["GPUs Used (key)"] = df["GPUs Used"].apply(lambda x: tuple(x) if isinstance(x, list) else tuple())
+
+    # 3) Drop rows missing essentials
+    essentials = numeric_cols + ["GPUs Used (key)"]
+    df_cleaned = df.dropna(subset=essentials).copy()
+
+    if df_cleaned.empty:
+        # Nothing to optimize; return empty frames shaped like callers expect
+        return df_cleaned, df_cleaned
+
+    # 4) Fastest runtime per (Stations, CPUs, GPUs, VRAM, Model) bucket
+    grp_cols = [
+        "Number of Stations Used",
+        "Number of CPUs Allocated for Ray to Use",
+        "GPUs Used (key)",
+        "Inference Actor Memory Limit (MB)",
+        "Model Used",
+    ]
+    idx = (
+        df_cleaned
+        .groupby(grp_cols)["Total Run time for Picker (s)"]
+        .idxmin()
+    )
+    optimal_concurrent_preds = df_cleaned.loc[idx].copy()
+
+    # For readability in outputs, show GPUs as list again
+    optimal_concurrent_preds["GPUs Used"] = optimal_concurrent_preds["GPUs Used (key)"].apply(list)
+    # (Optional) drop helper key in the returned table
+    optimal_concurrent_preds.drop(columns=["GPUs Used (key)"], inplace=True, errors="ignore")
+
+    # 5) “Moderate VRAM” window; if empty, fall back safely
+    vram_min = df_cleaned["Inference Actor Memory Limit (MB)"].quantile(0.25)
+    vram_max = df_cleaned["Inference Actor Memory Limit (MB)"].quantile(0.75)
+    df_moderate_vram = df_cleaned[
+        (df_cleaned["Inference Actor Memory Limit (MB)"] >= vram_min)
+        & (df_cleaned["Inference Actor Memory Limit (MB)"] <= vram_max)
+    ].copy()
+    if df_moderate_vram.empty:
+        df_moderate_vram = df_cleaned.copy()
+
+    # Highest stations first, then fastest runtime
+    best_overall_config = df_moderate_vram.sort_values(
+        by=["Number of Stations Used", "Total Run time for Picker (s)"],
+        ascending=[False, True],
+    ).iloc[0]
+
+    formatted_output = {
+        "Trial Number": best_overall_config["Trial Number"],
+        "Number of Stations Used": best_overall_config["Number of Stations Used"],
+        "Total Number of Timechunks": best_overall_config["Total Number of Timechunks"],
+        "Concurrent Timechunks Used": best_overall_config["Concurrent Timechunks Used"],
+        "Length of Timechunk (min)": str(best_overall_config["Length of Timechunk (min)"]),
+        "Total Waveform Analysis Timespace (min)": str(best_overall_config["Total Waveform Analysis Timespace (min)"]),
+        "Number of Concurrent Station Tasks per Timechunk": best_overall_config["Number of Concurrent Station Tasks"],
+        "Number of CPUs Allocated for Ray to Use": best_overall_config["Number of CPUs Allocated for Ray to Use"],
+        "GPUs Used": list(best_overall_config.get("GPUs Used (key)", ())) or best_overall_config.get("GPUs Used", []),
+        "Inference Actor Memory Limit (MB)": best_overall_config["Inference Actor Memory Limit (MB)"],
+        "Intra-parallelism Threads": best_overall_config["Intra-parallelism Threads"],
+        "Inter-parallelism Threads": best_overall_config["Inter-parallelism Threads"],
+        "Total Run time for Picker (s)": best_overall_config["Total Run time for Picker (s)"],
+        "Model Used": best_overall_config["Model Used"],
+        "Trial Success": best_overall_config["Trial Success"],
+        "Error Message": best_overall_config["Error Message"],
+    }
+    best_overall_df = pd.DataFrame([formatted_output])
+
+    return optimal_concurrent_preds, best_overall_df
+
+"""
+find_optimal_configuration_cpu/gpu returns back the best overall usecase results configuration, and takes those values to be used as the 
+current operation's runtime configuration."""
+def find_optimal_configuration_cpu(best_overall_usecase:bool, eval_sys_results_dir:str, cpu:int=None, station_count:int=None): 
+    # Check if eval_sys_results_dir is valid
+    if not eval_sys_results_dir or not os.path.isdir(eval_sys_results_dir):
+        print(f"Error: The provided directory path '{eval_sys_results_dir}' is invalid or does not exist.")
+        print("Please provide a valid directory path for the input parameter 'csv_dir'.")
+        return exit()  # Exit early if the directory is invalid
+    
+    if best_overall_usecase is True: 
+        file_path = f"{eval_sys_results_dir}/best_overall_usecase_cpu.csv"
+
+        # Check if the CSV file exists before reading
+        if not os.path.exists(file_path):
+            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
+            return exit()
+
+        # Load the CSV
+        df_best_overall = pd.read_csv(file_path)
+        # Convert into a dictionary for easy access
+        best_config_dict = df_best_overall.set_index(df_best_overall.columns[0]).to_dict()[df_best_overall.columns[1]]
+
+        # Extract required values
+        num_cpus = best_config_dict.get("Number of CPUs Allocated for Ray to Use")
+        waveform_timespace = best_config_dict.get("Total Waveform Analysis Timespace (min)")
+        total_num_timechunks = best_config_dict.get("Total Number of Timechunks")
+        num_concurrent_timechunks = best_config_dict.get("Concurrent Timechunks Used")
+        length_of_timechunks = best_config_dict.get("Length of Timechunk (min)")
+        num_concurrent_stations = best_config_dict.get("Number of Concurrent Station Tasks")
+        intra_threads = best_config_dict.get("Intra-parallelism Threads")
+        inter_threads = best_config_dict.get("Inter-parallelism Threads")
+        num_stations = best_config_dict.get("Number of Stations Used")
+        total_runtime = best_config_dict.get("Total Run time for Picker (s)")
+        model_used = best_config_dict.get("Model Used")
+        
+        print("\nBest Overall Usecase Configuration Based on Trial Data:")
+        print(f"Model Used: {model_used}\n"
+        f"CPU: {num_cpus}\n"
+        f"Intra-parallelism Threads: {intra_threads}\n"
+        f"Inter-parallelism Threads: {inter_threads}\n"
+        f"Waveform Timespace: {waveform_timespace}"
+        f"Total Number of Timechunks: {total_num_timechunks}"
+        f"Length of Timechunks (min): {length_of_timechunks}"
+        f"Concurrent Timechunks: {num_concurrent_stations}\n"
+        f"Concurrent Stations: {num_concurrent_stations}\n"
+        f"Stations: {num_stations}\n"
+        f"Total Runtime (s): {total_runtime}")
+
+        # Return the extracted values
+        return int(float(num_cpus)), int(float(num_concurrent_stations)), int(float(intra_threads)), int(float(inter_threads)), int(float(num_stations))
+    
+    else: # Optimal Configuration for User-Specified CPUs and Number of Stations to use
+        # Ensure valid CPU and station count values
+        if cpu is None or station_count is None:
+            print("Error: CPU and station_count must have valid values.")
+            return exit()
+        
+        file_path = f"{eval_sys_results_dir}/optimal_configurations_cpu.csv"
+
+        # Check if the CSV file exists before reading
+        if not os.path.exists(file_path):
+            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
+            return exit() 
+        
+        
+        df_optimal = pd.read_csv(file_path)
+
+        # Convert relevant columns to numeric, handling NaNs gracefully
+        df_optimal["Number of Stations Used"] = pd.to_numeric(df_optimal["Number of Stations Used"], errors="coerce")
+        df_optimal["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df_optimal["Number of CPUs Allocated for Ray to Use"], errors="coerce")
+        df_optimal["Number of Concurrent Station Tasks"] = pd.to_numeric(df_optimal["Number of Concurrent Station Tasks"], errors="coerce")
+        df_optimal["Total Run time for Picker (s)"] = pd.to_numeric(df_optimal["Total Run time for Picker (s)"], errors="coerce")
+        filtered_df = df_optimal[
+        (df_optimal["Number of CPUs Allocated for Ray to Use"] == cpu) &
+        (df_optimal["Number of Stations Used"] == station_count)]
+        if filtered_df.empty:
+            print("No matching configuration found. Please enter a valid entry.")
+            exit() 
+
+        # Find the best configuration (fastest runtime)
+        best_config = filtered_df.nsmallest(1, "Total Run time for Picker (s)").iloc[0]
+        
+        print("\nBest Configuration for Requested Input Parameters Based on Trial Data:")
+        print(f"Model Used: {best_config.get('Model Used')}\n"
+            f"CPU: {cpu}\nConcurrent Predictions: {best_config['Number of Concurrent Station Tasks']}\n"
+            f"Intra-parallelism Threads: {best_config['Intra-parallelism Threads']}\n"
+            f"Inter-parallelism Threads: {best_config['Inter-parallelism Threads']}\n"
+            f"Stations: {station_count}\nTotal Runtime (s): {best_config['Total Run time for Picker (s)']}")
+
+        return int(float(cpu)), int(float(best_config["Number of Concurrent Station Tasks"])), int(float(best_config["Intra-parallelism Threads"])), int(float(best_config["Inter-parallelism Threads"])), int(float(station_count))
+
+
+def find_optimal_configuration_gpu(best_overall_usecase: bool, eval_sys_results_dir: str, num_cpus: int = None, num_gpus: list = None, station_count: int = None):
+    """
+    Find the optimal GPU configuration for a given number of CPUs, GPUs, and stations.
+    Returns the best configuration including CPUs, concurrent predictions, intra/inter parallelism threads,
+    GPUs, VRAM, and stations.
+    """
+
+    # Check if eval_sys_results_dir is valid
+    if not eval_sys_results_dir or not os.path.isdir(eval_sys_results_dir):
+        print(f"Error: The provided directory path '{eval_sys_results_dir}' is invalid or does not exist.")
+        print("Please provide a valid directory path for the input parameter 'csv_dir'.")
+        return None  # Exit early if the directory is invalid
+
+    if best_overall_usecase:
+        file_path = f"{eval_sys_results_dir}/best_overall_usecase_gpu.csv"
+
+        # Check if the CSV file exists before reading
+        if not os.path.exists(file_path):
+            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
+            return None
+
+        # Load the CSV
+        df_best_overall = pd.read_csv(file_path, header=None, index_col=0)
+
+        # Convert into a dictionary for easy access
+        best_config_dict = df_best_overall.to_dict()[1]  # Extract key-value pairs
+
+        # Extract required values
+        num_cpus = best_config_dict.get("Number of CPUs Allocated for Ray to Use")
+        num_concurrent_stations = best_config_dict.get("Number of Concurrent Station Tasks")
+        intra_threads = best_config_dict.get("Intra-parallelism Threads")
+        inter_threads = best_config_dict.get("Inter-parallelism Threads")
+        num_stations = best_config_dict.get("Number of Stations Used")
+        total_runtime = best_config_dict.get("Total Run time for Picker (s)")
+        vram_used = best_config_dict.get("Inference Actor Memory Limit (MB)")
+        num_gpus_st = best_config_dict.get("GPUs Used")
+        num_gpus = ast.literal_eval(num_gpus_st)
+        model_used = best_config_dict.get("Model Used")
+        
+        print("\nBest Overall Usecase Configuration Based on Trial Data:")
+        print(f"Model Used: {model_used}\n"
+              f"CPU: {num_cpus}\n"
+              f"GPU ID(s): {num_gpus}\n"
+              f"Concurrent Predictions: {num_concurrent_stations}\n"
+              f"Intra-parallelism Threads: {intra_threads}\n"
+              f"Inter-parallelism Threads: {inter_threads}\n"
+              f"Stations: {num_stations}\n"
+              f"Inference Actor Memory Limit (MB): {vram_used}\n"
+              f"Total Runtime (s): {total_runtime}")
+
+        return int(float(num_cpus)), int(float(num_concurrent_stations)), int(float(intra_threads)), int(float(inter_threads)), num_gpus, int(float(vram_used)), int(float(num_stations))
+
+    else:  # Optimal Configuration for User-Specified CPUs, GPUs, and Number of Stations to use
+        # Ensure valid CPU, GPU, and station count values
+        if num_cpus is None or station_count is None or num_gpus is None:
+            print("Error: num_cpus, station_count, and num_gpus must have valid values.")
+            return None
+
+        file_path = f"{eval_sys_results_dir}/optimal_configurations_gpu.csv"
+
+        # Check if the CSV file exists before reading
+        if not os.path.exists(file_path):
+            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
+            return None
+
+        df_optimal = pd.read_csv(file_path)
+
+        # Convert relevant columns to numeric, handling NaNs gracefully
+        df_optimal["Number of Stations Used"] = pd.to_numeric(df_optimal["Number of Stations Used"], errors="coerce")
+        df_optimal["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df_optimal["Number of CPUs Allocated for Ray to Use"], errors="coerce")
+        df_optimal["Number of Concurrent Station Tasks"] = pd.to_numeric(df_optimal["Number of Concurrent Station Tasks"], errors="coerce")
+        df_optimal["Total Run time for Picker (s)"] = pd.to_numeric(df_optimal["Total Run time for Picker (s)"], errors="coerce")
+        df_optimal["Inference Actor Memory Limit (MB)"] = pd.to_numeric(df_optimal["Inference Actor Memory Limit (MB)"], errors="coerce")
+
+        # Convert "GPUs Used" from string representation to list
+        df_optimal["GPUs Used"] = df_optimal["GPUs Used"].apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else x)
+
+        # Convert GPU lists to tuples for comparison
+        df_optimal["GPUs Used"] = df_optimal["GPUs Used"].apply(lambda x: tuple(x) if isinstance(x, list) else (x,))
+
+        # Ensure num_gpus is in tuple format for comparison
+        num_gpus_tuple = tuple(num_gpus) if isinstance(num_gpus, list) else (num_gpus,)
+
+        filtered_df = df_optimal[
+            (df_optimal["Number of CPUs Allocated for Ray to Use"] == num_cpus) &
+            (df_optimal["GPUs Used"] == num_gpus_tuple) &
+            (df_optimal["Number of Stations Used"] == station_count)
+        ]
+
+        if filtered_df.empty:
+            print("No matching configuration found. Please enter a valid entry.")
+            exit()
+
+        # Find the best configuration (fastest runtime)
+        best_config = filtered_df.nsmallest(1, "Total Run time for Picker (s)").iloc[0]
+        
+        print("\nBest Configuration for Requested Application Usecase Based on Trial Data:")
+        print(f"Model Used: {best_config.get('Model Used')}\n"
+              f"CPU: {num_cpus}\n"
+              f"GPU: {num_gpus}\n"
+              f"Concurrent Predictions: {best_config['Number of Concurrent Station Tasks']}\n"
+              f"Intra-parallelism Threads: {best_config['Intra-parallelism Threads']}\n"
+              f"Inter-parallelism Threads: {best_config['Inter-parallelism Threads']}\n"
+              f"Stations: {station_count}\n"
+              f"Inference Actor Memory Limit (MB): {best_config['Inference Actor Memory Limit (MB)']}\n"
+              f"Total Runtime (s): {best_config['Total Run time for Picker (s)']}")
+
+        return int(float(best_config["Number of CPUs Allocated for Ray to Use"])), \
+               int(float(best_config["Number of Concurrent Station Tasks"])), \
+               int(float(best_config["Intra-parallelism Threads"])), \
+               int(float(best_config["Inter-parallelism Threads"])), \
+               num_gpus, \
+               int(float(best_config["Inference Actor Memory Limit (MB)"])), \
+               int(float(station_count))
+