compute_mlm.sh

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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$SCRIPT_DIR"

START=400
END=900
STEP=100

LOG="runs/tmp2.log"
TMUX_LOG_DIR="runs/tmux_logs"
mkdir -p "$(dirname "$LOG")" "$TMUX_LOG_DIR"

for x in $(seq "$START" "$STEP" "$END"); do
  y=$((x + 99))
  idx=$(((x - START) / STEP))
  gpu=$((idx % 8))
  session="mlm_${x}_${y}_gpu${gpu}"
  sess_log="${TMUX_LOG_DIR}/${session}.log"

  if tmux has-session -t "$session" 2>/dev/null; then
    echo "tmux session exists, skipping: $session"
    continue
  fi

  tmux new-session -d -s "$session" "bash --noprofile --norc -c '
    set -eo pipefail

    # log everything
    mkdir -p \"$(printf %q "$TMUX_LOG_DIR")\"
    exec > >(tee -a \"$(printf %q "$sess_log")\") 2>&1

    echo \"[$session] START \$(date)\"

    # keep pane open on failure for debugging
    trap '\''echo \"[$session] ERROR at \$(date)\"; exec bash'\'' ERR

    cd \"$(printf %q "$REPO_ROOT")\"

    if [ -x \"$(printf %q "$REPO_ROOT")/.venv/bin/python\" ]; then
      PYTHON_BIN=\"$(printf %q "$REPO_ROOT")/.venv/bin/python\"
    elif [ -x /mnt/promoter-pretrain/.venv/bin/python ]; then
      PYTHON_BIN=/mnt/promoter-pretrain/.venv/bin/python
    else
      # bring conda into this clean shell as a fallback
      if [ -f /opt/conda/etc/profile.d/conda.sh ]; then
        source /opt/conda/etc/profile.d/conda.sh
      elif [ -f \"\$HOME/miniconda3/etc/profile.d/conda.sh\" ]; then
        source \"\$HOME/miniconda3/etc/profile.d/conda.sh\"
      else
        echo \"No workspace venv found and conda.sh not found (checked /opt/conda and ~/miniconda3)\" >&2
        exit 1
      fi

      conda activate bert24
      PYTHON_BIN=python
    fi

    export HF_HOME=/mnt/promoter-pretrain/.cache/hf
    mkdir -p \"\$HF_HOME\"
    set -euo pipefail

    for i in {${x}..${y}}; do
      date
      echo \"$session: \$i --> \$(date)\" >> \"$(printf %q "$LOG")\"
      CUDA_VISIBLE_DEVICES=${gpu} \"\$PYTHON_BIN\" src/data/fast_compute_MLM_metrics.py \
        --job_number \"\${i}\" \
        --batch_size 96 \
        --num_workers 8 \
        --mlm_efficiency_path /dev/shm/mlm_outputs/ \
        --append_mlm_efficiency \
        --data_dir \"$(printf %q "$REPO_ROOT")/data/promoters_s3\"
    done

    echo \"[$session] DONE \$(date)\"
    exec bash
  '"

  echo "Started: $session (GPU $gpu, jobs ${x}..${y})"
done

echo "tmux ls"
echo "logs: $TMUX_LOG_DIR/"