run_build_validation_dir_4runs_modal.sh

#!/usr/bin/env bash
# Build the PRM "validation" dataset (with branch-off vectors) across 4
# translation directions, 15 kernels each, using Modal cloud GPUs for
# model inference.  All output goes to data/validation_dir/.
#
# Model inference (Coconut translate + continuations) → Modal remote GPUs
# HeCBench validation → skipped (--skip-validation, scores set to 0)
#
# Usage:
#   bash run_build_validation_dir_4runs_modal.sh
#   bash run_build_validation_dir_4runs_modal.sh --resume
#
# Override config or GPU spec:
#   COCONUT_CONFIG=code_translation_70b.yaml MODAL_GPU="H200:3" bash ...
#
# The 4 directions, in order:
#   1. cuda   -> omp
#   2. omp    -> cuda
#   3. serial -> cuda
#   4. serial -> omp
#
set -o pipefail

source ~/miniconda3/etc/profile.d/conda.sh
conda activate coconut
set -u

export COCONUT_CONFIG="${COCONUT_CONFIG:-../coconut_large_model_inference/args/code_translation_70b.yaml}"
export PYTHONUNBUFFERED=1
# Keep this enabled for PRM data; continuation forwards remain batched.
SAVE_CONTINUATION_VECTORS="${SAVE_CONTINUATION_VECTORS:-1}"

RETRY_EXIT_CODE=33
RETRY_DELAY_SECS=10
MAX_ATTEMPTS=1
LIMIT=15
OUTPUT_DIR=data/validation_dir
RUN_TS="$(date -u +%Y%m%dT%H%M%SZ)"
LOG_DIR="logs/validation_dir_${RUN_TS}"
mkdir -p "${LOG_DIR}"

# Modal GPU spec — override with MODAL_GPU env var
MODAL_GPU="${MODAL_GPU:-A100-80GB:3}" #H200:2}" #A100-80GB:3}"

# Compute sensible per-GPU memory cap from the config path
if [[ "$MODAL_GPU" == *":"* ]]; then
    NUM_GPUS="${MODAL_GPU##*:}"
else
    NUM_GPUS=1
fi
if [[ "$COCONUT_CONFIG" == *"7b.yaml" ]]; then
    TOTAL_MEM=14
elif [[ "$COCONUT_CONFIG" == *"70b.yaml" ]]; then
    TOTAL_MEM=136
elif [[ "$COCONUT_CONFIG" == *"30b.yaml" ]]; then
    TOTAL_MEM=60
elif [[ "$COCONUT_CONFIG" == *"14b.yaml" ]]; then
    TOTAL_MEM=30
else
    TOTAL_MEM=136
fi
MEM_PER_GPU=$(( (TOTAL_MEM / NUM_GPUS) + 1 ))

echo "[modal-wrapper] Config:          ${COCONUT_CONFIG}"
echo "[modal-wrapper] Modal GPU:       ${MODAL_GPU} (${NUM_GPUS} GPU(s)) -> max_memory_per_gpu=${MEM_PER_GPU}GiB"
echo "[modal-wrapper] Output dir:      ${OUTPUT_DIR}"
echo "[modal-wrapper] Limit/direction: ${LIMIT}"
echo "[modal-wrapper] Validation:      SKIP (--skip-validation)"
echo "[modal-wrapper] Save cont vecs:  ${SAVE_CONTINUATION_VECTORS}"

# ---------------------------------------------------------------------------
# Run one direction with the same retry/back-off loop.
# Args: <from_api> <to_api>
# ---------------------------------------------------------------------------
run_one_direction() {
    local FROM_API="$1"
    local TO_API="$2"
    local DIR_LIMIT="${3:-$LIMIT}"
    local EXTRA_ARGS=("${@:4}")
    local LOG_FILE="${LOG_DIR}/build_${FROM_API}_${TO_API}.log"
    local CONT_VECTOR_ARGS=()

    if [[ "${SAVE_CONTINUATION_VECTORS}" == "1" ]]; then
        CONT_VECTOR_ARGS+=(--save-continuation-vectors)
    fi

    echo ""
    echo "=========================================================="
    echo "[modal-wrapper] Direction: ${FROM_API} -> ${TO_API}"
    echo "=========================================================="
    echo "[modal-wrapper] ${FROM_API}->${TO_API} log: ${LOG_FILE}"

    local attempt=1
    while [[ $attempt -le $MAX_ATTEMPTS ]]; do
        echo "[modal-wrapper] ${FROM_API}->${TO_API} attempt ${attempt}/${MAX_ATTEMPTS}"

        local BATCH_SIZE
        if [[ $attempt -ge 4 ]]; then
            BATCH_SIZE=7
        else
            BATCH_SIZE=13
        fi

        set +e
        python scripts/run_build_dataset.py \
            --config "$COCONUT_CONFIG" \
            --split train \
            --log-level DEBUG \
            --log-file "${LOG_FILE}" \
            --from-api "$FROM_API" --to-api "$TO_API" \
            -n "$DIR_LIMIT" \
            --max-workers 8 \
            --num-resamples 2 \
            --max-gen-batch-size ${BATCH_SIZE} \
            --max-memory-per-gpu "${MEM_PER_GPU}GiB" \
            --output-dir "$OUTPUT_DIR" \
            --modal \
            --modal-gpu "${MODAL_GPU}" \
            --resume \
            --give-up \
            --skip-validation \
            --skip-kernel-names d2q9-bgk dp dxtc2 \
            "${CONT_VECTOR_ARGS[@]}" \
            "${EXTRA_ARGS[@]}"
        local rc=$?
        set -e

        if [[ $rc -eq 0 ]]; then
            echo "[modal-wrapper] ${FROM_API}->${TO_API} completed successfully"
            return 0
        fi

        if [[ $rc -eq $RETRY_EXIT_CODE ]]; then
            echo "[modal-wrapper] exited with ${RETRY_EXIT_CODE}; retrying in ${RETRY_DELAY_SECS}s"
            sleep "$RETRY_DELAY_SECS"
            attempt=$((attempt + 1))
            continue
        fi

        echo "[modal-wrapper] ${FROM_API}->${TO_API} failed with non-retryable exit code ${rc}"
        return "$rc"
    done

    echo "[modal-wrapper] ${FROM_API}->${TO_API} exhausted ${MAX_ATTEMPTS} attempts"
    return 1
}

# ---------------------------------------------------------------------------
# 4 directions sequentially.
# ---------------------------------------------------------------------------
pids=()
labels=()

run_one_direction cuda   omp    30 "$@" &
pids+=("$!")
labels+=("cuda->omp")

run_one_direction omp    cuda   30 "$@" &
pids+=("$!")
labels+=("omp->cuda")

run_one_direction serial cuda   30 "$@" &
pids+=("$!")
labels+=("serial->cuda")

run_one_direction serial omp    30 "$@" &
pids+=("$!")
labels+=("serial->omp")

overall_rc=0
for idx in "${!pids[@]}"; do
    pid="${pids[$idx]}"
    label="${labels[$idx]}"
    if wait "$pid"; then
        echo "[modal-wrapper] ${label} finished successfully"
    else
        rc=$?
        echo "[modal-wrapper] ${label} failed with exit code ${rc}"
        if [[ $overall_rc -eq 0 ]]; then
            overall_rc=$rc
        fi
    fi
done

if [[ $overall_rc -ne 0 ]]; then
    echo ""
    echo "[modal-wrapper] One or more directions failed. Logs in ${LOG_DIR}"
    exit "$overall_rc"
fi

echo ""
echo "[modal-wrapper] All 4 directions completed. Output in ${OUTPUT_DIR}"
echo "[modal-wrapper] Per-direction logs in ${LOG_DIR}"