name: livekit-wakeword-train

resources:
  cloud: nebius
  accelerators: L40S:1
  cpus: 8+
  memory: 32+
  disk_size: 500
  autostop:
    idle_minutes: 10
    down: false

workdir: .

file_mounts:
  /outputs:
    source: nebius://wakeword-model-bucket
    mode: MOUNT

setup: |
  sudo apt-get update && sudo apt-get install -y espeak-ng libsndfile1 ffmpeg sox portaudio19-dev
  curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.cargo/bin:$PATH"
  # Persist data and output outside workdir so they survive re-syncs
  mkdir -p ${HOME_DIR}/${DATA_DIR} ${HOME_DIR}/${OUTPUT_DIR}
  ln -sfn ${HOME_DIR}/${DATA_DIR} data
  ln -sfn ${HOME_DIR}/${OUTPUT_DIR} output
  uv sync --all-extras
  uv run livekit-wakeword setup --config configs/${CONFIG_FILE}

run: |
  export PATH="$HOME/.cargo/bin:$PATH"
  ln -sfn ${HOME_DIR}/${DATA_DIR} data
  ln -sfn ${HOME_DIR}/${OUTPUT_DIR} output
  uv run livekit-wakeword run configs/${CONFIG_FILE}
  MODEL_DIR="${HOME_DIR}/${OUTPUT_DIR}/${MODEL_NAME}"
  ARTIFACT_DIR="/outputs/${MODEL_NAME}"
  mkdir -p "$ARTIFACT_DIR"
  cp "$MODEL_DIR/${MODEL_NAME}.pt" "$ARTIFACT_DIR/"
  cp "$MODEL_DIR/${MODEL_NAME}.onnx" "$ARTIFACT_DIR/"
  cp "$MODEL_DIR/${MODEL_NAME}_metrics.json" "$ARTIFACT_DIR/"
  cp "$MODEL_DIR/${MODEL_NAME}_det.png" "$ARTIFACT_DIR/"
  cp "$MODEL_DIR/${MODEL_NAME}_eval.json" "$ARTIFACT_DIR/"

envs:
  CONFIG_FILE: prod.yaml
  MODEL_NAME: hey_livekit
  HOME_DIR: /home/ubuntu
  DATA_DIR: persistent/data
  OUTPUT_DIR: persistent/output