team-rexx / training /cloud_as_run /launch_script.sh

Team Rexx system: inference.py, final predictions (Codabench 145), training code

5318caf verified 18 days ago

1.94 kB

#!/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/"

	#!/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/"