#!/usr/bin/env bash
# Start one Spot A100 Vertex AI job that runs a Python script with vLLM available (see script_job.sh).
# Usage: scripts/vertex/launch_script.sh <run_id> <local script.py> <input file>... -- [script args]
# Env: REGION, GPUS. Inputs land in ./inputs/ inside the job; write outputs to /gcs/ai-erp-fiji-sinhalammlu/jobs/<run_id>/out/.
# Example: scripts/vertex/launch_script.sh rw1 scripts/train/rewrite_traces.py data/train/cd2_pre.jsonl -- \
#            inputs/cd2_pre.jsonl /gcs/ai-erp-fiji-sinhalammlu/jobs/rw1/out/cd2_rw.jsonl
set -euo pipefail
cd "$(dirname "$0")/../.."
PROJECT=ai-erp-fiji
REGION=${REGION:-us-central1}
GPUS=${GPUS:-1}   # 2 for the 26B/31B teachers
BUCKET=gs://ai-erp-fiji-sinhalammlu
IMAGE=vllm/vllm-openai:v0.32.0-cu129

RUN=$1; SCRIPT=$2; shift 2
INPUTS=()
while [[ $# -gt 0 && $1 != "--" ]]; do INPUTS+=("$1"); shift; done
[[ ${1:-} == "--" ]] && shift
ARGS=("$@")

gcloud storage cp "$SCRIPT" scripts/vertex/script_job.sh "$BUCKET/code/" --quiet
gcloud storage cp "${INPUTS[@]}" "$BUCKET/jobs/$RUN/inputs/" --quiet

mkdir -p experiments/jobs
CFG=experiments/jobs/${RUN}.yaml
{
  echo "workerPoolSpecs:"
  echo "  - machineSpec: {machineType: a2-highgpu-${GPUS}g, acceleratorType: NVIDIA_TESLA_A100, acceleratorCount: $GPUS}"
  echo "    replicaCount: 1"
  echo "    diskSpec: {bootDiskType: pd-ssd, bootDiskSizeGb: 200}"
  echo "    containerSpec:"
  echo "      imageUri: $IMAGE"
  echo "      command: [bash, /gcs/ai-erp-fiji-sinhalammlu/code/script_job.sh]"
  printf '      args: ["%s", "%s"' "$RUN" "$(basename "$SCRIPT")"
  for a in "${ARGS[@]+"${ARGS[@]}"}"; do printf ', "%s"' "$a"; done
  echo "]"
  echo "scheduling: {strategy: SPOT, timeout: 14400s}"
} > "$CFG"
gcloud ai custom-jobs create --project=$PROJECT --region=$REGION \
  --display-name="script-${RUN}-${REGION}" --config="$CFG" --format="value(name)"
echo "log: $BUCKET/jobs/$RUN/job.log   out: $BUCKET/jobs/$RUN/out/"