social-data-attribution-w2/scripts/attribution/trackstar_common.sh

configure_variant() {
    QUERY_DIR="${QUERY_DIR:-queries/$VARIANT}"
    case "$VARIANT" in
        base | base_olmobaseeval)
            MODEL="allenai/Olmo-3-1025-7B"
            ;;
        instruct_base)
            MODEL="allenai/Olmo-3-7B-Instruct"
            ;;
        instruct_cot)
            MODEL="allenai/Olmo-3-7B-Instruct"
            ;;
        *)
            echo "ERROR: invalid variant: $VARIANT"
            exit 1
            ;;
    esac
}


resolve_run_id() {
    case "$START_PHASE" in
        reduce | build)
            if [ -z "$RUN_ID" ]; then
                RUN_ID="$(date -u +%Y%m%dT%H%M%SZ)_$$"
            fi
            ;;
        score | aggregate | dot_score)
            if [ -z "$RUN_ID" ]; then
                echo "ERROR: --run-id is required when --phase $START_PHASE"
                exit 1
            fi
            ;;
        *)
            echo "ERROR: invalid phase: $START_PHASE"
            exit 1
            ;;
    esac
}


build_sbatch_flags() {
    SBATCH_FLAGS=()
    if [ -n "${SLURM_ACCOUNT:-}" ]; then SBATCH_FLAGS+=(--account="$SLURM_ACCOUNT"); fi
    if [ -n "${SLURM_PARTITION:-}" ]; then SBATCH_FLAGS+=(--partition="$SLURM_PARTITION"); fi
    if [ -n "${SLURM_QOS:-}" ]; then SBATCH_FLAGS+=(--qos="$SLURM_QOS"); fi
    if [ -n "${SLURM_CONSTRAINT:-}" ]; then SBATCH_FLAGS+=(-C "$SLURM_CONSTRAINT"); fi
    if [ -n "${CPUS_PER_TASK:-}" ]; then SBATCH_FLAGS+=(--cpus-per-task="$CPUS_PER_TASK"); fi
    if [ -n "${MEM:-}" ]; then SBATCH_FLAGS+=(--mem="$MEM"); fi
    if [ -n "${GPU_WALLTIME:-}" ]; then SBATCH_FLAGS+=(--time="$GPU_WALLTIME"); fi
}


run_sbatch() {
    local desc="$1"
    shift

    if [ "$DRY_RUN" = true ]; then
        printf >&2 "[dry-run] sbatch"
        printf >&2 " %q" "$@"
        printf >&2 "\n"
        printf "DRY_RUN_JOB_%s_%s\n" "$$" "$RANDOM"
        return
    fi

    local job_id
    job_id=$(sbatch --parsable "$@")
    printf >&2 "Submitted %s: job %s\n" "$desc" "$job_id"
    printf "%s\n" "$job_id"
}


load_query_names() {
    local query_names=()
    local jsonl

    if [ "$START_PHASE" = "reduce" ]; then
        if [ ! -d "$QUERY_DIR" ]; then
            echo "ERROR: Query directory not found: $QUERY_DIR"
            echo "Run data-attribution-trackstar-query --variant $VARIANT first."
            exit 1
        fi
        for jsonl in "$QUERY_DIR"/*.jsonl; do
            [ -s "$jsonl" ] || continue
            query_names+=("queries_$(basename "$jsonl" .jsonl | sed 's/^olmes_//')")
        done
    else
        local reduce_path
        if [ ! -d "$REDUCE_DIR" ]; then
            echo "ERROR: Reduce directory not found: $REDUCE_DIR"
            echo "Run Phase 1 (reduce) first or pass the correct --run-id."
            exit 1
        fi
        for reduce_path in "$REDUCE_DIR"/queries_*; do
            [ -d "$reduce_path" ] || continue
            query_names+=("$(basename "$reduce_path")")
        done
    fi

    if [ ${#query_names[@]} -eq 0 ]; then
        echo "ERROR: No query inputs found for variant $VARIANT"
        exit 1
    fi

    printf "%s\n" "${query_names[@]}"
}


print_launch_header() {
    echo "=============================================="
    echo "  TrackStar Pipeline Launch"
    echo "=============================================="
    echo "Mode:         $MODE"
    echo "Variant:      $VARIANT"
    echo "Model:        $MODEL"
    echo "Run ID:       $RUN_ID"
    echo "Run root:     $RUN_ROOT"
    echo "Parallelism:  $PARALLELISM"
    echo "Start phase:  $START_PHASE"
    echo "Top-K:        $TOP_K"
    echo "Token batch:  ${TOKEN_BATCH:-default}"
    echo "Precond dir:  ${PRECONDITIONER_DIR:-none}"
    echo "SBATCH flags: ${SBATCH_FLAGS[*]:-none}"
    echo "=============================================="
}


print_launch_footer() {
    echo ""
    echo "=============================================="
    echo "  Pipeline submitted (Mode B)"
    echo "=============================================="
    echo "Run ID:        $RUN_ID"
    echo "Reduce jobs:   ${#REDUCE_JOB_IDS[@]}"
    echo "Score jobs:    ${#SCORE_JOB_IDS[@]}"
    echo "Aggregate job: $AGG_JOB_ID"
    echo ""
    echo "Monitor:"
    ALL_IDS=("$AGG_JOB_ID")
    if [ ${#SCORE_JOB_IDS[@]} -gt 0 ]; then ALL_IDS=("${SCORE_JOB_IDS[@]}" "${ALL_IDS[@]}"); fi
    if [ ${#REDUCE_JOB_IDS[@]} -gt 0 ]; then ALL_IDS=("${REDUCE_JOB_IDS[@]}" "${ALL_IDS[@]}"); fi
    echo "  squeue -j $(IFS=,; echo "${ALL_IDS[*]}") -o '%i %j %t %M %R'"
}


print_mode_a_footer() {
    echo ""
    echo "=============================================="
    echo "  Pipeline submitted (Mode A)"
    echo "=============================================="
    echo "Run ID:          $RUN_ID"
    echo "Build jobs:      ${#BUILD_JOB_IDS[@]}"
    echo "Query build jobs: ${#QUERY_BUILD_JOB_IDS[@]}"
    echo "Dot score job:   $DOT_SCORE_JOB_ID"
    echo ""
    echo "Monitor:"
    ALL_IDS=("$DOT_SCORE_JOB_ID")
    if [ ${#QUERY_BUILD_JOB_IDS[@]} -gt 0 ]; then ALL_IDS=("${QUERY_BUILD_JOB_IDS[@]}" "${ALL_IDS[@]}"); fi
    if [ ${#BUILD_JOB_IDS[@]} -gt 0 ]; then ALL_IDS=("${BUILD_JOB_IDS[@]}" "${ALL_IDS[@]}"); fi
    echo "  squeue -j $(IFS=,; echo "${ALL_IDS[*]}") -o '%i %j %t %M %R'"
}