Update compile_flash.sh

compile_flash.sh

@@ -1,50 +1,70 @@
 #!/bin/bash
 set -e
-
-echo "🔧 [Aduc-SDR] Compilador Persistente de FlashAttention (GPU Mode)"
-echo "----------------------------------------------------------------"
+echo "🔧 Compilador FlashAttention — modo GPU (configurado para infra grande)"
 
 WHEEL_DIR="/data/wheel"
-mkdir -p "$WHEEL_DIR"
+TMP_DIR="/data/tmp"
+mkdir -p "$WHEEL_DIR" "$TMP_DIR"
+export TMPDIR="$TMP_DIR"
 
 FLASH_WHL=$(find "$WHEEL_DIR" -name "flash_attn-*.whl" | head -n 1)
-
 if [[ -f "$FLASH_WHL" ]]; then
-    echo "📦 Wheel já existente: $FLASH_WHL"
-    echo "✅ Reutilizando compilação anterior."
-    pip install --no-cache-dir "$FLASH_WHL"
-    exit 0
+  echo "📦 Wheel já existe: $FLASH_WHL"
+  pip install --no-cache-dir "$FLASH_WHL"
+  exit 0
 fi
 
-echo "⚙️ Nenhum wheel encontrado — iniciando compilação com CUDA (modo seguro)..."
-
-# --- Controle de paralelismo e memória ---
-export MAX_JOBS=4                # usa metade das CPUs
-export OMP_NUM_THREADS=4
-export CUDA_VISIBLE_DEVICES="0"  # usa apenas 1 GPU
-export TORCH_CUDA_ARCH_LIST="8.9"
-export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:256"
+echo "⚙️ Nenhum wheel encontrado — iniciando compilação controlada com CUDA..."
 
-# --- Flags de build CUDA ---
+# ----- Ajustes seguros para ambiente grande (8x L40S) -----
+#export MAX_JOBS=${MAX_JOBS:-128}          # jobs de compilação (ninja/cmake)
+#export OMP_NUM_THREADS=${OMP_NUM_THREADS:-8}
+#export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}  # usar apenas 1 GPU para build
+#export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-"8.9"}
+#export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-"max_split_size_mb:512"}
 export USE_FLASH_ATTENTION_WITH_CUDA=1
 export FORCE_CUDA=1
 
-# --- Compilação controlada ---
-echo "🚧 Compilando flash-attn==2.8.3 com suporte CUDA..."
-pip wheel flash-attn==2.8.3 --no-build-isolation -w "$WHEEL_DIR" || {
-    echo "⚠️ Falha na compilação CUDA. Tentando fallback CPU-only..."
-    export USE_FLASH_ATTENTION_WITH_CUDA=0
-    export FORCE_CUDA=0
-    pip wheel flash-attn==2.8.3 --no-build-isolation -w "$WHEEL_DIR"
-}
+echo "   > MAX_JOBS=$MAX_JOBS, OMP_NUM_THREADS=$OMP_NUM_THREADS, CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+echo "   > TMPDIR=$TMPDIR, WHEEL_DIR=$WHEEL_DIR"
+
+# Tenta build com CUDA
+set -o pipefail
+if pip wheel flash-attn==2.8.0 --no-build-isolation -w "$WHEEL_DIR" 2>&1 | tee /tmp/flash_build.log; then
+  echo "✅ Build com CUDA concluído."
+else
+  echo "⚠️ Build CUDA falhou; tentando fallback CPU-only (mais lento). Verifique /tmp/flash_build.log"
+  export USE_FLASH_ATTENTION_WITH_CUDA=0
+  export FORCE_CUDA=0
+  pip wheel flash-attn==2.8.0 --no-build-isolation -w "$WHEEL_DIR"
+fi
 
-# --- Instala o wheel gerado ---
 FLASH_WHL=$(find "$WHEEL_DIR" -name "flash_attn-*.whl" | head -n 1)
 if [[ -f "$FLASH_WHL" ]]; then
-    echo "✅ Wheel gerado com sucesso: $FLASH_WHL"
-    pip install --no-cache-dir "$FLASH_WHL"
-    echo "📁 Salvo em: $WHEEL_DIR"
+  echo "✅ Wheel gerado: $FLASH_WHL"
+  pip install --no-cache-dir "$FLASH_WHL"
+  echo "📁 Wheel salvo em $WHEEL_DIR"
+else
+  echo "❌ Erro: não foi possível gerar o wheel. Ver logs em /tmp/flash_build.log"
+  exit 1
+fi
+
+# ----- Upload opcional para Hugging Face -----
+if [ "${HF_UPLOAD_WHEELS:-0}" = "1" ]; then
+  echo "⬆️ Upload habilitado. Verificando HF_TOKEN..."
+  python3 - <<'PY'
+import os
+from huggingface_hub import HfApi, HfFolder
+repo = os.environ.get("SELF_HF_REPO_ID","caarleexx/Flash")
+token = os.getenv("HF_TOKEN") or HfFolder.get_token()
+if not token:
+    raise SystemExit("HF_TOKEN ausente; upload desabilitado")
+api = HfApi(token=token)
+api.upload_folder(folder_path="/data/wheel", repo_id=repo, repo_type="model",
+                  allow_patterns=["*.whl","NGC-DL-CONTAINER-LICENSE"],
+                  ignore_patterns=["**/src/**","**/*.log","**/logs/**",".git/**"])
+print("Upload concluído.")
+PY
 else
-    echo "❌ Falha ao gerar wheel FlashAttention."
-    exit 1
+  echo "ℹ️ Upload HF desabilitado (HF_UPLOAD_WHEELS!=1)"
 fi