File size: 5,045 Bytes
315e4cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
# Shared helpers. Source after REPO / EVIE_ROOT are set.
evie_resolve_python() {
  if [[ -z "${PYTHON:-}" ]]; then
    if [[ -n "${VIRTUAL_ENV:-}" && -x "$VIRTUAL_ENV/bin/python" ]]; then
      PYTHON="$VIRTUAL_ENV/bin/python"
    else
      PYTHON="$(command -v python3)"
    fi
  fi
  [[ -x "$PYTHON" ]] || { echo "[fatal] set PYTHON or VIRTUAL_ENV"; exit 2; }
}

evie_pythonpath() {
  export PYTHONPATH="$REPO/colpali:$REPO/code/shared${PYTHONPATH:+:$PYTHONPATH}"
}

evie_real() {
  local p="$1"
  if command -v realpath >/dev/null 2>&1; then
    realpath -m "$p" 2>/dev/null || echo "$p"
  else
    echo "$p"
  fi
}

# Refuse dumps inside the venv (the old $EVIE_ROOT/env footgun).
evie_forbid_venv_path() {
  local path="$1" label="$2"
  local child parent roots=() p
  child="$(evie_real "$path")"
  [[ -n "${VIRTUAL_ENV:-}" ]] && roots+=("$(evie_real "$VIRTUAL_ENV")")
  roots+=("$(evie_real "$EVIE_ROOT/env")")
  roots+=("$(evie_real "$EVIE_ROOT/venv")")
  roots+=("$(evie_real "$EVIE_ROOT/.venv")")
  roots+=("$(evie_real "$EVIE_ROOT/train-env")")
  if [[ -n "${PYTHON:-}" && -x "$PYTHON" ]]; then
    parent="$(cd "$(dirname "$PYTHON")/.." && pwd)"
    [[ -f "$parent/pyvenv.cfg" ]] && roots+=("$(evie_real "$parent")")
  fi
  for p in "${roots[@]}"; do
    [[ -n "$p" && "$p" != "/" ]] || continue
    if [[ "$child" == "$p" || "$child" == "$p"/* ]]; then
      echo "[fatal] $label cannot sit inside the Python env ($p): $path" >&2
      echo "[fatal] checkpoints -> \$RUNS_DIR (\$EVIE_ROOT/runs); caches -> \$EVIE_ROOT/.cache" >&2
      exit 2
    fi
  done
}

# Work product roots. All gitignored. Never under env/.
evie_workdirs() {
  export RUNS_DIR="${RUNS_DIR:-$EVIE_ROOT/runs}"
  export LOG_DIR="${LOG_DIR:-$EVIE_ROOT/logs}"
  export EVAL_ROOT="${EVAL_ROOT:-$EVIE_ROOT/data}"
  export EVIE_TMP="${EVIE_TMP:-$EVIE_ROOT/.tmp}"
  export CACHE_DIR="${CACHE_DIR:-$EVIE_ROOT/.cache}"
  export HF_HOME="${HF_HOME:-$CACHE_DIR/huggingface}"
  export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-$HF_HOME/datasets}"
  export HF_HUB_CACHE="${HF_HUB_CACHE:-$HF_HOME/hub}"
  export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-$HF_HUB_CACHE}"
  export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-$HF_HUB_CACHE}"
  export TORCH_HOME="${TORCH_HOME:-$CACHE_DIR/torch}"
  evie_forbid_venv_path "$RUNS_DIR" "RUNS_DIR"
  evie_forbid_venv_path "$LOG_DIR" "LOG_DIR"
  evie_forbid_venv_path "$EVAL_ROOT" "EVAL_ROOT"
  evie_forbid_venv_path "$CACHE_DIR" "CACHE_DIR"
  evie_forbid_venv_path "$HF_HOME" "HF_HOME"
  mkdir -p "$RUNS_DIR" "$LOG_DIR" "$EVIE_TMP" "$CACHE_DIR"
}

evie_nccl() {
  evie_workdirs
  export TOKENIZERS_PARALLELISM=false
  export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
  export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}"
  export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC="${TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC:-480}"
  export TORCH_NCCL_ENABLE_MONITORING="${TORCH_NCCL_ENABLE_MONITORING:-1}"
  export TORCH_NCCL_ASYNC_ERROR_HANDLING="${TORCH_NCCL_ASYNC_ERROR_HANDLING:-1}"
  export TORCH_FR_BUFFER_SIZE="${TORCH_FR_BUFFER_SIZE:-2000}"
  export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-0}"
}

evie_batch() {
  NPROC_PER_NODE="${NPROC_PER_NODE:-$("$PYTHON" -c 'import torch; print(torch.cuda.device_count())')}"
  TOTAL_GPUS=$((NNODES * NPROC_PER_NODE))
  MICRO_GLOBAL=$((TOTAL_GPUS * BSZ))
  GRAD_ACCUM="${GRAD_ACCUM:-$((EFF_BATCH / MICRO_GLOBAL))}"
  ACTUAL_GLOBAL=$((MICRO_GLOBAL * GRAD_ACCUM))
}

evie_wandb() {
  if [[ ",$REPORT_TO," != *",wandb,"* ]]; then
    return
  fi
  export WANDB_PROJECT WANDB_NAME WANDB_DIR
  if [[ -z "${WANDB_API_KEY:-}" ]]; then
    echo "[warn] WANDB_API_KEY absent; wandb offline"
    export WANDB_MODE=offline
  else
    export WANDB_MODE="${WANDB_MODE:-online}"
  fi
}

# Rank 0 prepares OUTPUT_DIR; other nodes wait on a shared ready token.
evie_prepare_output() {
  local token="${MASTER_ADDR}_${MASTER_PORT}_${RUN_NAME}"
  local ready="$EVIE_TMP/launch_${RUN_NAME}.ready"
  evie_forbid_venv_path "$OUTPUT_DIR" "OUTPUT_DIR"
  mkdir -p "$EVIE_TMP" "$RUNS_DIR"
  if [[ "$NODE_RANK" == "0" ]]; then
    rm -f "$ready"
    if [[ -d "$OUTPUT_DIR" && -n "$(ls -A "$OUTPUT_DIR" 2>/dev/null || true)" && -z "${RESUME:-}" ]]; then
      if [[ "${OVERWRITE:-auto}" == "0" ]]; then
        echo "[fatal] non-empty output: $OUTPUT_DIR (OVERWRITE=auto/1)"
        printf 'ERROR:%s\n' "$token" > "$ready"
        exit 2
      fi
      local archive="${OUTPUT_DIR}_archive_$(date +%Y%m%d_%H%M%S)"
      mv "$OUTPUT_DIR" "$archive"
      echo "[output] archived -> $archive"
    fi
    mkdir -p "$OUTPUT_DIR"
    printf '%s\n' "$token" > "$ready"
  else
    local ready_val=""
    for _ in $(seq 1 300); do
      ready_val=""
      IFS= read -r ready_val < "$ready" || true
      [[ "$ready_val" == "$token" ]] && break
      [[ "$ready_val" == "ERROR:$token" ]] && { echo "[fatal] rank0 rejected output setup"; exit 2; }
      sleep 1
    done
    [[ "$ready_val" == "$token" ]] || { echo "[fatal] timed out waiting for rank0 output setup"; exit 2; }
  fi
  export EVIE_OUTPUT_PREPARED=1
}