Update builder.sh

builder.sh

@@ -1,13 +1,22 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-
-
-
 echo "======================================================="
 echo " ADUC-SDR — Builder (Triton + Apex + Q8 + FlashAttention + LayerNorm)"
 echo "======================================================="
 
+# ---------------------- Toolchain p/ compilar ----------------------
+if command -v apt-get >/dev/null 2>&1; then
+  apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
+    build-essential clang lld cmake ninja-build git python3-dev \
+    && rm -rf /var/lib/apt/lists/*
+fi
+export CC="${CC:-clang}"
+export CXX="${CXX:-clang++}"
+export LD="${LD:-lld}"
+
+chmod -R 777 /app/wheels /app/whells || true
+
 # ---------------------- Caches persistentes ----------------------
 if [ -d /data ]; then
   export HF_HOME="${HF_HOME:-/data/.cache/huggingface}"
@@ -35,12 +44,10 @@ export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.9}"
 export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
 export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
 
-# ---------------------- Dependências de build ----------------------
-
-
+# ---------------------- Base Python deps ----------------------
+python -m pip install -U pip setuptools wheel build cmake ninja packaging "huggingface_hub[hf_transfer]" || true
 
-
-# Tags do ambiente
+# ---------------------- Identificadores de ambiente ----------------------
 PY_TAG="$(python -c 'import sys; print(f"cp{sys.version_info[0]}{sys.version_info[1]}")' 2>/dev/null || echo cp310)"
 TORCH_VER="$(python - <<'PY'
 try:
@@ -101,9 +108,7 @@ PY
 }
 
 hf_upload_wheels_if_available() {
-  # Faz upload de wheels geradas para o repo HF (opcional)
-  # Requer HF_TOKEN configurado (env)
-  local pkg_prefix="$1" # ex: triton-
+  local pkg_prefix="$1"
   if [ -z "${HF_TOKEN:-}" ]; then
     echo "ℹ️ HF_TOKEN não definido; skip upload wheels ${pkg_prefix}"
     return 0
@@ -116,10 +121,10 @@ hf_upload_wheels_if_available() {
   fi
   echo "☁️ Enviando wheels ${pkg_prefix}* para HF repo ${SELF_HF_REPO_ID} ..."
   python - <<'PY'
-import os
+import os, glob
 from huggingface_hub import HfApi, HfFolder
-repo=os.environ.get("SELF_HF_REPO_ID","carlex3321/aduc-sdr")
-token=os.getenv("HF_TOKEN") or HfFolder.get_token()
+repo = os.environ.get("SELF_HF_REPO_ID","carlex3321/aduc-sdr")
+token = os.environ.get("HF_TOKEN") or HfFolder.get_token()
 if not token:
   raise SystemExit(0)
 api=HfApi(token=token)
@@ -183,37 +188,20 @@ raise SystemExit(0 if ok else 1)
 PY
 }
 
-check_flash_layer_norm_wheel() {
-python - <<'PY'
-ok=False
-try:
-  import importlib
-  # tenta símbolos típicos do LN nativo
-  for name in ("flash_attn.ops.layer_norm","flash_attn.layers.layer_norm"):
-    try:
-      m=importlib.import_module(name)
-      ok=True; break
-    except Exception:
-      pass
-except Exception:
-  ok=False
-raise SystemExit(0 if ok else 1)
-PY
-}
-
 # ---------------------- Builders ----------------------
+TRITON_REPO="${TRITON_REPO:-https://github.com/openai/triton.git}"
+TRITON_REF="${TRITON_REF:-main}"
+
 build_triton() {
   echo "🔧 build_triton — preparando"
   rm -rf /app/wheels/src/triton
-  git clone --depth 1 https://github.com/openai/triton.git /app/wheels/src/triton
-  # Sugeridos (opcional):
-  export TRITON_BUILD_WITH_CLANG_LLD=true
-  export TRITON_BUILD_WITH_CCACHE=true
-  export TRITON_HOME=/app/.triton
+  git clone --depth 1 --branch "$TRITON_REF" "$TRITON_REPO" /app/wheels/src/triton
   pushd /app/wheels/src/triton >/dev/null
+  # Tentar wheel offline (gera cache reusável); se falhar, instala editable
   pip wheel --no-build-isolation --no-deps . -w /app/wheels || true
   popd >/dev/null
-  use_wheel_or_build "triton-*.whl" "/app/wheels/src/triton" "pip install --no-build-isolation /app/wheels/src/triton"
+  use_wheel_or_build "triton-*.whl" "/app/wheels/src/triton" "pip install --no-build-isolation -e /app/wheels/src/triton || pip install --no-build-isolation /app/wheels/src/triton"
+  # Upload opcional (se wheel existir)
   hf_upload_wheels_if_available "triton-"
   echo "✅ build_triton — OK"
 }
@@ -258,10 +246,6 @@ build_flash () {
   rm -rf "$SRC"
   git clone --depth 1 --branch "$FLASH_ATTENTION_TAG" https://github.com/Dao-AILab/flash-attention.git "$SRC"
 
-  export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.9}"
-  export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
-  export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
-
   # 1) Pacote principal
   python -m pip wheel --no-build-isolation --no-deps "$SRC" -w /app/wheels || true
   local W_MAIN="$(ls -t /app/wheels/flash_attn-*.whl 2>/dev/null | head -n1 || true)"
@@ -271,24 +255,28 @@ build_flash () {
     python -m pip install --no-build-isolation "$SRC" || true
   fi
 
-  # 2) Submódulo attn/csrc/layer_norm — empacotar e instalar
-  echo "🔧 flash-attn layer_norm — build wheel de csrc/layer_norm"
-  local LN_SRC="$SRC/attn/csrc/layer_norm"
+  # 2) Submódulo attn/csrc/layer_norm — gerar e instalar wheel
+  echo "🔧 flash-attn layer_norm — build wheel"
+  local LN_SRC="$SRC/csrc/layer_norm"
+  if [ ! -d "$LN_SRC" ]; then
+    LN_SRC="$SRC/attn/csrc/layer_norm"
+  fi
   if [ -d "$LN_SRC" ]; then
-    if [ -f "$LN_SRC/setup.py" ] || [ -f "$LN_SRC/pyproject.toml" ]; then
-      python -m pip wheel --no-build-isolation --no-deps "$LN_SRC" -w /app/wheels || true
-      local W_LN="$(ls -t /app/wheels/*layer*norm*whl /app/wheels/flash_attn_layer_norm-*.whl 2>/dev/null | head -n1 || true)"
-      if [ -n "${W_LN}" ]; then
-        python -m pip install -U --no-deps "${W_LN}" || true
-      else
-        pip install --no-build-isolation "$LN_SRC" || true
-      fi
+    pushd "$LN_SRC" >/dev/null
+    pip wheel --no-build-isolation --no-deps . -w /app/wheels || pip wheel --use-pep517 --no-deps . -w /app/wheels || true
+    popd >/dev/null
+    local W_LN="$(ls -t /app/wheels/dropout_layer_norm-*.whl 2>/dev/null | head -n1 || true)"
+    if [ -n "${W_LN}" ]; then
+      python -m pip install -U --no-deps "${W_LN}" || true
     else
-      # fallback simples se não houver metadados de build locais
-      pip install --no-build-isolation "$LN_SRC" || true
+      echo "⚠️ Wheel dropout_layer_norm não encontrada, instalando do source (PEP 517 se possível)"
+      pushd "$LN_SRC" >/dev/null
+      pip install --use-pep517 --no-build-isolation . || pip install .
+      popd >/dev/null
     fi
+    hf_upload_wheels_if_available "dropout_layer_norm-"
   else
-    echo "ℹ️ Diretório $LN_SRC não encontrado na tag $FLASH_ATTENTION_TAG; seguindo."
+    echo "ℹ️ Diretório de layer_norm não encontrado (tag ${FLASH_ATTENTION_TAG})."
   fi
 
   hf_upload_wheels_if_available "flash_attn-"
@@ -322,26 +310,15 @@ ensure_pkg () {
 }
 
 # ---------------------- Execução ----------------------
-# Ordem recomendada: Triton -> Apex -> Q8 -> FlashAttention (+ layer_norm) -> BitsAndBytes
-#ensure_pkg "triton"       check_triton 
-#build_triton
-#|| true
-#ensure_pkg "apex"         check_apex   build_apex   || true
-#ensure_pkg "q8_kernels"   check_q8     build_q8     || true
-#ensure_pkg "flash_attn"   check_flash  build_flash  || true
+build_triton
+ensure_pkg "triton"       check_triton build_triton || true
+ensure_pkg "apex"         check_apex   build_apex   || true
+ensure_pkg "q8_kernels"   check_q8     build_q8     || true
+ensure_pkg "flash_attn"   check_flash  build_flash  || true
 
 # opcional
-build_triton
-build_apex
-build_flash
-build_q8
 build_bitsandbytes
 
-check_triton
-check_apex
-check_q8
-check_flash
-
 echo "======================================================="
 echo " Build concluído. Resumo de módulos:"
 python - <<'PY'
@@ -358,6 +335,11 @@ try:
   print(" - apex.normalization: OK")
 except Exception as e:
   print(" - apex.normalization: FAIL ->", e)
+try:
+  import dropout_layer_norm
+  print(" - dropout_layer_norm: OK")
+except Exception as e:
+  print(" - dropout_layer_norm: FAIL ->", e)
 PY
 echo "======================================================="
 echo " ✔️ Builder finalizado."