#!/usr/bin/env bash
set -euo pipefail

project=/home/miaochengyu/starVLA-infoflow-codex
runtime=/mnt/miaochengyu/starvla-infoflow-runtime
adapted="${project}/external_baselines/frameskip_adapted"
upstream="${project}/external_baselines/upstream_readonly/FrameSkip-965becf63e9cd5cc339fbf7fde40d1db47efa67a/starVLA"
config="${adapted}/configs/florence_bridge_frameskip_gaussian_r02_seed117465.yaml"
required="${adapted}/cache/required_draws.json"
visual="${adapted}/assets/hf_home/hub/models--timm--vit_small_patch14_dinov2.lvd142m/snapshots/4610ca143709d58a633b6397a74412c2c3842454/model.safetensors"
output="${adapted}/cache/shards"
logs="${adapted}/logs/cache_build"
workers_per_gpu=4
shard_count=$((workers_per_gpu * 2))

mkdir -p "${output}" "${logs}"
cd "${project}"
export PYTHONPATH="${project}:${PYTHONPATH:-}"
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=4
export MKL_NUM_THREADS=4
export OPENBLAS_NUM_THREADS=4
export NUMEXPR_NUM_THREADS=4

for shard in $(seq 0 $((shard_count - 1))); do
  if (( shard < workers_per_gpu )); then gpu=2; else gpu=3; fi
  CUDA_VISIBLE_DEVICES="${gpu}" "${runtime}/.venv/bin/python" \
    "${adapted}/build_cache_shard.py" \
    --project-root "${project}" \
    --upstream-root "${upstream}" \
    --config "${config}" \
    --required-draws "${required}" \
    --visual-checkpoint "${visual}" \
    --output-dir "${output}" \
    --shard-index "${shard}" \
    --shard-count "${shard_count}" \
    >"${logs}/shard_${shard}.log" 2>&1 &
  printf '%s\n' "$!" >"${logs}/shard_${shard}.pid"
done
wait