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

if [ $# -ne 1 ]; then
  echo "Usage: bash launch_and_log.sh <train_script>"
  exit 1
fi

TRAIN_SCRIPT="$1"
RUN_ID="$(date +%Y%m%d_%H%M%S)"
RUN_DIR="/raid/data/weifeng/Pre_train/swift_runs/runs/${RUN_ID}"

if command -v python >/dev/null 2>&1; then
  PYTHON_BIN=python
elif command -v python3 >/dev/null 2>&1; then
  PYTHON_BIN=python3
else
  echo "Error: neither python nor python3 was found in PATH."
  exit 1
fi

mkdir -p "${RUN_DIR}"

cp "${TRAIN_SCRIPT}" "${RUN_DIR}/command.sh"
env | sort > "${RUN_DIR}/env.txt"
nvidia-smi > "${RUN_DIR}/nvidia-smi.txt"
"${PYTHON_BIN}" -m pip freeze > "${RUN_DIR}/pip_freeze.txt"

echo "Run directory: ${RUN_DIR}"
echo "Logging to: ${RUN_DIR}/train.log"

bash "${TRAIN_SCRIPT}" 2>&1 | tee "${RUN_DIR}/train.log"