scripts/wait_and_run_cosmos.sh

#!/usr/bin/env bash
# Wait for a GPU with enough free memory, then run Cosmos-Transfer2.5 on it.
#
# This is a shared host. Cosmos-Transfer2.5-2B needs 65.4 GB for a single
# control and, measured here, ~85 GB in multicontrol mode (it loads all four
# control branches even when only two carry weight). Launching into
# insufficient memory does not just fail our run -- it raises memory pressure
# on whatever else is on that GPU, so this waits for real headroom instead of
# gambling.
#
# Usage:
#   scripts/wait_and_run_cosmos.sh <spec.json> <out_dir> <need_gb> [extra args...]
#
# Environment: HF_TOKEN required. POLL_SECONDS (default 300).

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
NVSHIM="$REPO_ROOT/.nvshim"
POLL="${POLL_SECONDS:-300}"

if [[ $# -lt 3 ]]; then
    echo "usage: $0 <spec.json> <out_dir> <need_gb> [extra args...]" >&2
    exit 2
fi
SPEC="$1"; OUT="$2"; NEED_GB="$3"; shift 3

echo "waiting for a GPU with >= ${NEED_GB} GB free (polling every ${POLL}s)"

while true; do
    # nvidia-smi only works with the shim on this host; see run_cosmos_transfer.sh.
    best_idx=-1; best_free=0
    while IFS=, read -r idx free; do
        idx="${idx// /}"; free="${free// /}"; free="${free//MiB/}"
        [[ -z "$free" ]] && continue
        if (( free > best_free )); then best_free=$free; best_idx=$idx; fi
    done < <(LD_LIBRARY_PATH="$NVSHIM" nvidia-smi \
                --query-gpu=index,memory.free --format=csv,noheader 2>/dev/null || true)

    best_gb=$(( best_free / 1024 ))
    echo "$(date '+%H:%M:%S')  best: GPU${best_idx} with ${best_gb} GB free (need ${NEED_GB})"

    if (( best_idx >= 0 && best_gb >= NEED_GB )); then
        echo "launching on GPU${best_idx}"
        COSMOS_GPUS="$best_idx" exec "$REPO_ROOT/scripts/run_cosmos_transfer.sh" \
            "$SPEC" "$OUT" "$@"
    fi
    sleep "$POLL"
done