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

#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
cd "$REPO_DIR"
source "$SCRIPT_DIR/trackstar_common.sh"
source "$SCRIPT_DIR/trackstar_phases.sh"

VARIANT=""
PARALLELISM=8
MODE="mode_a"
START_PHASE=""
TOP_K=100
RUN_ID=""
DRY_RUN=false
RUN_ROOT="${RUN_ROOT:-runs/trackstar}"
PRECONDITIONER_DIR=""
TOKEN_BATCH="${TOKEN_BATCH:-}"

while [ $# -gt 0 ]; do
    case "$1" in
        --variant) VARIANT="$2"; shift 2 ;;
        --mode) MODE="$2"; shift 2 ;;
        --parallelism) PARALLELISM="$2"; shift 2 ;;
        --phase) START_PHASE="$2"; shift 2 ;;
        --top-k) TOP_K="$2"; shift 2 ;;
        --run-id) RUN_ID="$2"; shift 2 ;;
        --run-root) RUN_ROOT="$2"; shift 2 ;;
        --preconditioner-dir) PRECONDITIONER_DIR="$2"; shift 2 ;;
        --token-batch) TOKEN_BATCH="$2"; shift 2 ;;
        --dry-run) DRY_RUN=true; shift ;;
        *) echo "Unknown argument: $1"; exit 1 ;;
    esac
done

if [ -z "$VARIANT" ]; then
    echo "ERROR: --variant is required (base | base_olmobaseeval | instruct_base | instruct_cot)"
    exit 1
fi

if [ -z "$START_PHASE" ]; then
    if [ "$MODE" = "mode_a" ]; then
        START_PHASE="build"
    else
        START_PHASE="reduce"
    fi
fi

configure_variant
resolve_run_id
build_sbatch_flags

EXPORT_BASE="ALL"
EXPORT_BASE+=",PRECISION=${PRECISION:-fp32}"
if [ -n "$TOKEN_BATCH" ]; then
    EXPORT_BASE+=",TOKEN_BATCH=$TOKEN_BATCH"
fi
SHARD_DIR="${SHARD_DIR:-$RUN_ROOT/shards}"

print_launch_header

if [ "$MODE" = "mode_a" ]; then
    BUILD_DIR="$RUN_ROOT/builds/$VARIANT/$RUN_ID"
    QUERY_BUILD_DIR="$RUN_ROOT/query_builds/$VARIANT/$RUN_ID"
    BUILD_JOB_IDS=()
    QUERY_BUILD_JOB_IDS=()

    if [ "$START_PHASE" = "build" ]; then
        submit_build_jobs
        submit_query_build_jobs
    fi
    submit_dot_score_job
    print_mode_a_footer
else
    REDUCE_DIR="$RUN_ROOT/reduces/$VARIANT/$RUN_ID"
    SCORE_DIR="$RUN_ROOT/scores/$VARIANT/$RUN_ID"
    REDUCE_JOB_IDS=()
    SCORE_JOB_IDS=()

    if [ "$START_PHASE" = "reduce" ]; then
        submit_reduce_jobs
    fi
    if [ "$START_PHASE" = "reduce" ] || [ "$START_PHASE" = "score" ]; then
        submit_score_jobs
    fi
    submit_aggregate_job
    print_launch_footer
fi