wattameter.utils package

wattameter.utils.postprocessing module

wattameter.utils.postprocessing.align_and_concat_df(_list_df, dt=None, start_at_0=False)

Create a single dataframe from multiple dataframes, aligning them in time.

The column labels of the input dataframes are prefixed with their index in the input list to avoid collisions.

Parameters:
  • _list_df – List of pandas DataFrames to combine. Each DataFrame must have a DateTimeIndex.

  • dt – Time step in seconds for the new index. If None, the average time step across all dataframes is used. (default: None)

  • start_at_0 – If True, reset the index to start at 0 seconds. If False, use time stamps that start at a common start time among the dataframes. (default: False)

Returns:

A single pandas DataFrame with aligned time index and combined data.

wattameter.utils.postprocessing.file_to_df(f, timestamp_fmt='%Y-%m-%d_%H:%M:%S.%f', header=None, skip_lines=1)

Convert an output file from Wattameter Tracker to a pandas DataFrame.

Parameters:
  • f – Open file object to read from.

  • timestamp_fmt – Format string for parsing timestamps. (default: '%Y-%m-%d_%H:%M:%S.%f')

  • header – List of column names. If None, the header is read from the file. (default: None)

  • skip_lines – Number of lines to skip at the beginning of the file. (default: 1)

wattameter.utils.wattameter module

#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: 2025, Alliance for Sustainable Energy, LLC
#
# This script is used to run the WattAMeter CLI tool.
# It captures the output and PID of the process, allowing for graceful termination on timeout.

# Usage function to display help
usage() {
    echo "Usage: $0 [-i|--index run_id] [-s|--suffix suffix] [-q|--quiet] [wattameter-options]"
    echo "-i, --id    run_id  : Specify a run identifier for this WattAMeter instance"
    echo "-s, --suffix suffix : Specify a custom suffix for log file naming"
    echo "-q, --quiet         : Quiet mode; suppress startup messages"
    echo "-h, --help          : Display this help message"
    echo "wattameter-options  : Additional options to pass to the wattameter command"
    echo ""
    echo "Note: Default suffix is hostname."
    echo "      If --id is provided, default suffix becomes run_id-hostname."
    exit 0
}

main() {
    # Get the hostname of the current node
    local NODE=$(hostname)

    # Parse arguments manually
    local RUN_ID=""
    local SUFFIX=${NODE}
    local CUSTOM_SUFFIX=false
    local QUIET=false
    local EXTRA_ARGS=()
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -i|--id) RUN_ID="$2"; shift 2 ;;
            -s|--suffix) SUFFIX="$2"; CUSTOM_SUFFIX=true; shift 2 ;;
            -q|--quiet) QUIET=true; shift ;;
            -h|--help) usage ;;
            *) EXTRA_ARGS+=("$1"); shift ;;
        esac
    done
    set -- "${EXTRA_ARGS[@]}"  # Restore positional parameters

    # Determine suffix based on flags
    if [[ "$CUSTOM_SUFFIX" = false && -n "${RUN_ID}" ]]; then
        SUFFIX="${RUN_ID}-${NODE}"
    fi

    # Set log file name
    local log_file="wattameter-${SUFFIX}.txt"
    if [[ "${QUIET}" = false ]]; then
        echo "Logging execution on ${NODE} to ${log_file}"
    fi

    # Build wattameter command arguments
    local WATTAMETER_ARGS="--suffix ${SUFFIX}"
    if [ -n "${RUN_ID}" ]; then
        WATTAMETER_ARGS="${WATTAMETER_ARGS} --id ${RUN_ID}"
    fi

    # Start the tracking and log the output
    wattameter ${WATTAMETER_ARGS} "$@" > "${log_file}" 2>&1 &
    local WATTAMETER_PID=$!

    # Gracefully terminates the tracking process on exit.
    local SIGNAL=""
    on_exit() {
        local TIMESTAMP=$(date +"%Y-%m-%d %H:%M:%S")
        echo "${TIMESTAMP}: WattAMeter interrupted on ${NODE} by signal ${SIGNAL}. Terminating..."

        # Interrupt the WattAMeter process
        kill -INT "$WATTAMETER_PID" 2>/dev/null
        wait "$WATTAMETER_PID" 2>/dev/null
        while kill -0 "$WATTAMETER_PID" 2>/dev/null; do
            sleep 0.1
        done

        local TIMESTAMP=$(date +"%Y-%m-%d %H:%M:%S")
        echo "${TIMESTAMP}: WattAMeter has been terminated on node ${NODE}."
    }
    trap 'SIGNAL=INT; on_exit' INT
    trap 'SIGNAL=TERM; on_exit' TERM
    trap 'SIGNAL=HUP; on_exit' HUP

    # Wait for the WattAMeter process to finish
    wait "$WATTAMETER_PID"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi

wattameter.utils.wattawait module

#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: 2025, Alliance for Sustainable Energy, LLC
#
# This script is used to wait until "run ID" can be read from the file path.
# If no file path is given, it will use the wattameter_powerlog_filename
# utility to get the file path for the current node.

# Usage function to display help
usage() {
    echo "Usage: $0 [-q] [-f filepath] ID"
    echo "  -q          Quiet mode (suppress output)"
    echo "  -f filepath Optional file path to monitor"
    echo "  ID          Run ID to wait for"
    exit "${1:-0}"
}

# Parse options
QUIET=false
FILEPATH=""
while getopts "qf:h" opt; do
    case "$opt" in
        q) QUIET=true ;;
        f) FILEPATH="$OPTARG" ;;
        h) usage 0 ;;
        *) usage 1 ;;
    esac
done
shift $((OPTIND - 1))

# Get the ID from the remaining arguments
if [ $# -ge 1 ]; then
    ID="$1"
else
    usage
fi

# If no filepath was provided via -f flag, generate it
if [ -z "$FILEPATH" ]; then
    # Get the WattAMeter powerlog file path for the current node
    NODE=$(hostname)
    FILEPATH=$(wattameter_powerlog_filename --suffix "${ID}-${NODE}")
fi

# Wait until ID can be read from the file
if [ "$QUIET" = false ]; then
    echo "Waiting for ${FILEPATH} to be ready for run ID ${ID}..."
fi
until grep -qs "run $ID" "${FILEPATH}"; do
    sleep 1  # Wait for 1 second before checking again
done

if [ "$QUIET" = false ]; then
    echo "${FILEPATH} is ready for run ID ${ID}."
fi
exit 0

wattameter.utils.slurm module

#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: 2025, Alliance for Sustainable Energy, LLC
#
# This script contains utility functions for SLURM job management.

start_wattameter () {
    # Error out if not running inside a SLURM job
    if [ -z "$SLURM_JOB_ID" ]; then
        echo "Error: start_wattameter must be run inside a SLURM job allocation."
        return 1
    fi

    # Set default value for SLURM_JOB_NUM_NODES if not defined
    if [ -z "$SLURM_JOB_NUM_NODES" ]; then
        local SLURM_JOB_NUM_NODES=1
    fi

    # Initialize global variables for wattameter steps
    if [ -z "$WATTAMETER_N_STARTED_STEPS" ]; then
        WATTAMETER_N_STARTED_STEPS=0
    fi
    if [ -z "$WATTAMETER_SLURM_STEP_IDS" ]; then
        WATTAMETER_SLURM_STEP_IDS=()
    fi

    # Use a suffix to differentiate multiple wattameter runs
    if [[ "$WATTAMETER_N_STARTED_STEPS" -gt 0 ]]; then
        local ID="$SLURM_JOB_ID-$WATTAMETER_N_STARTED_STEPS"
    else
        local ID="$SLURM_JOB_ID"
    fi
    WATTAMETER_N_STARTED_STEPS=$((WATTAMETER_N_STARTED_STEPS + 1))

    # Get the path of the wattameter script
    local WATTAPATH=$(python -c 'import wattameter; import os; print(os.path.dirname(wattameter.__file__))')
    local WATTASCRIPT="${WATTAPATH}/utils/wattameter.sh"
    local WATTAWAIT="${WATTAPATH}/utils/wattawait.sh"

    # Create sentinel to track wattameter start
    srun --overlap --wait=0 --nodes="$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 \
        "${WATTAWAIT}" -q "$ID" &
    local WAIT_PID=$!

    # Run wattameter on all nodes
    srun --overlap --wait=0 \
        --output="slurm-$ID-wattameter.txt" \
        --nodes="$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 \
        "${WATTASCRIPT}" -i "$ID" "$@" 2>/dev/null &

    # Wait for wattameter to start
    wait $WAIT_PID 2>/dev/null

    # Get the step ID from the last wattameter srun command
    local SANITY_CHECK=0
    while true; do
        local STEP_ID=$(squeue -j "$SLURM_JOB_ID" -h -s --format="%i %.13j" | grep "wattameter.sh$" | tail -1 | awk '{print $1}')
        if [ -n "$STEP_ID" ]; then
            break
        fi
        sleep 0.1
        SANITY_CHECK=$((SANITY_CHECK + 1))
        if [ $SANITY_CHECK -gt 600 ]; then
            echo "Error: Unable to retrieve wattameter SLURM step ID."
            return 1
        fi
    done
    WATTAMETER_SLURM_STEP_IDS+=("$STEP_ID")

    local TIMESTAMP=$(date +"%Y-%m-%d %H:%M:%S")
    echo "${TIMESTAMP}: Started WattAmeter job step $STEP_ID"
}

stop_wattameter () {    
    # Error out if not running inside a SLURM job
    if [ -z "$SLURM_JOB_ID" ]; then
        echo "Error: stop_wattameter must be run inside a SLURM job allocation."
        return 1
    fi

    # Cancel the last started wattameter step
    if [[ ${#WATTAMETER_SLURM_STEP_IDS[@]} -gt 0 ]]; then
        # Pop last wattameter STEP ID
        local STEP_ID="${WATTAMETER_SLURM_STEP_IDS[-1]}"
        unset 'WATTAMETER_SLURM_STEP_IDS[-1]'
        
        # Stop ID using scancel
        if [ -n "$STEP_ID" ]; then
            scancel --signal=INT $STEP_ID 2>/dev/null

            # Wait for the step to terminate
            local SANITY_CHECK=0
            while squeue -j "$SLURM_JOB_ID" -h -s --format "%i" | grep -q "^$STEP_ID$"; do
                sleep 0.1
                SANITY_CHECK=$((SANITY_CHECK + 1))
                if [ $SANITY_CHECK -gt 600 ]; then
                    echo "Warning: WattAMeter step $STEP_ID did not terminate in a timely manner."
                    break
                fi
            done

            local TIMESTAMP=$(date +"%Y-%m-%d %H:%M:%S")
            echo "${TIMESTAMP}: Stopped WattAmeter job step $STEP_ID"
        fi
    fi
}