Xxxxxxxxxxxxxxxxxxxxxx

builder.sh

@@ -1,10 +1,10 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-echo "🚀 Builder (Apex + Q8) — roda em runtime com GPU visível"
+echo "🚀 Builder (FlashAttn LayerNorm + Apex + Q8) — runtime com GPU visível"
 
 # ===== Config e diretórios =====
-export SELF_HF_REPO_ID="${SELF_HF_REPO_ID:-euIaxs22/Aduc-sdr}"   # Model repo no HF com wheels
+export SELF_HF_REPO_ID="${SELF_HF_REPO_ID:-euIaxs22/Aduc-sdr}"  # Repo no HF para wheels
 export HF_HOME="${HF_HOME:-/app/model_cache}"
 export HF_HUB_CACHE="${HF_HUB_CACHE:-$HF_HOME/hub}"
 export TORCH_HOME="${TORCH_HOME:-$HF_HOME/torch}"
@@ -15,6 +15,7 @@ mkdir -p /app/wheels /app/cuda_cache "$HF_HOME" "$TORCH_HOME" /app/wheels/src
 chmod -R 777 /app/wheels || true
 export CUDA_CACHE_PATH="/app/cuda_cache"
 
+# Preserve licença NGC em wheels (se presente no container base)
 if [ -f "/NGC-DL-CONTAINER-LICENSE" ]; then
   cp -f /NGC-DL-CONTAINER-LICENSE /app/wheels/NGC-DL-CONTAINER-LICENSE || true
 fi
@@ -44,14 +45,25 @@ PY
 )"
 echo "[env] PY_TAG=${PY_TAG} TORCH_VER=${TORCH_VER} CU_TAG=${CU_TAG}"
 
+# ============================================================================
+#                               CHECKERS
+# ============================================================================
 
-git clone https://github.com/Dao-AILab/flash-attention  
-cd flash-attention/csrc/layer_norm && MAX_JOBS=90 pip install .
-
-
+check_flashln () {
+python - <<'PY'
+import importlib.util, pkgutil
+cands = [
+    "flash_attn_layer_norm",     # empacotamentos recentes
+    "dropout_layer_norm",        # nome histórico do módulo C++
+    "flash_attn.ops.layer_norm", # caminho dentro do pacote flash_attn
+]
+ok = any(importlib.util.find_spec(n) is not None for n in cands)
+ok = ok or any(("flash" in m.name and "norm" in m.name) for m in pkgutil.iter_modules())
+raise SystemExit(0 if ok else 1)
+PY
+}
 
-# ===== Checkers =====
-check_apex() {
+check_apex () {
 python - <<'PY'
 try:
     from apex.normalization import FusedLayerNorm, FusedRMSNorm
@@ -63,7 +75,7 @@ raise SystemExit(0 if ok else 1)
 PY
 }
 
-check_q8() {
+check_q8 () {
 python - <<'PY'
 import importlib.util
 spec = importlib.util.find_spec("ltx_q8_kernels") or importlib.util.find_spec("q8_kernels")
@@ -71,15 +83,54 @@ raise SystemExit(0 if spec else 1)
 PY
 }
 
-# ===== Download do Hub =====
-install_from_hf () {
-  local PKG="$1"    # 'apex' ou 'q8_kernels'
-  echo "[hub] Verificando wheel de ${PKG} no repositório ${SELF_HF_REPO_ID}"
-  python - "$PKG" "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
-import os, sys
+# ============================================================================
+#                         DOWNLOAD DO HUB (GENÉRICO)
+# ============================================================================
+
+# Instala uma wheel do HF por padrão de nome (prefixo) e tags.
+# Uso: install_from_hf_by_prefix <prefixo>
+install_from_hf_by_prefix () {
+  local PREFIX="$1"
+  echo "[hub] Procurando wheels '${PREFIX}-*.whl' em ${SELF_HF_REPO_ID} com tags ${PY_TAG}/${CU_TAG}"
+  python - "$PREFIX" "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
+import os, sys, re
+from huggingface_hub import HfApi, hf_hub_download, HfFolder
+
+prefix, py_tag, cu_tag = sys.argv[1], sys.argv[2], sys.argv[3]
+repo = os.environ.get("SELF_HF_REPO_ID","euIaxs22/Aduc-sdr")
+api = HfApi(token=os.getenv("HF_TOKEN") or HfFolder.get_token())
+try:
+    files = api.list_repo_files(repo_id=repo, repo_type="model")
+except Exception:
+    raise SystemExit(0)
+
+def match_tags(name: str) -> bool:
+    # Requer Python tag; CUDA tag preferencial se existir na wheel
+    if py_tag not in name:
+        return False
+    return True
+
+cands = [f for f in files if f.endswith(".whl") and f.rsplit("/",1)[-1].startswith(prefix + "-") and match_tags(f)]
+pref = [f for f in cands if cu_tag and cu_tag in f] or cands
+if not pref:
+    raise SystemExit(0)
+
+# Seleciona o "maior" por ordenação lexicográfica como heurística de versão/compat
+target = sorted(pref, reverse=True)[0]
+print(target)
+path = hf_hub_download(repo_id=repo, filename=target, repo_type="model", local_dir="/app/wheels")
+print(path)
+PY
+}
+
+# Instala uma wheel do HF por padrão flexível para FlashAttn LayerNorm (aceita várias variantes de nome)
+install_flashln_from_hf () {
+  echo "[hub] Procurando wheels de FlashAttention LayerNorm em ${SELF_HF_REPO_ID}"
+  python - "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
+import os, sys, re
 from huggingface_hub import HfApi, hf_hub_download, HfFolder
 
-pkg, py_tag, cu_tag = sys.argv[1], sys.argv[2], sys.argv[3]
+py_tag, cu_tag = sys.argv[1], sys.argv[2]
 repo = os.environ.get("SELF_HF_REPO_ID","euIaxs22/Aduc-sdr")
 api = HfApi(token=os.getenv("HF_TOKEN") or HfFolder.get_token())
 try:
@@ -87,8 +138,19 @@ try:
 except Exception:
     raise SystemExit(0)
 
-cands = [f for f in files if f.endswith(".whl") and f.rsplit("/",1)[-1].startswith(pkg+"-") and py_tag in f]
-pref  = [f for f in cands if cu_tag and cu_tag in f] or cands
+def ok(fn: str) -> bool:
+    name = fn.rsplit("/",1)[-1]
+    if not name.endswith(".whl"): return False
+    if py_tag not in name: return False
+    # Padrões comuns de empacotamento
+    pats = [
+        r"^flash[_-]?attn[_-]?.*layer[_-]?norm-",
+        r"^dropout[_-]?layer[_-]?norm-",
+    ]
+    return any(re.search(p, name, flags=re.I) for p in pats)
+
+cands = [f for f in files if ok(f)]
+pref = [f for f in cands if cu_tag and cu_tag in f] or cands
 if not pref:
     raise SystemExit(0)
 
@@ -99,7 +161,61 @@ print(path)
 PY
 }
 
-# ===== Builders =====
+# ============================================================================
+#                                BUILDERS
+# ============================================================================
+
+build_flashln () {
+  local SRC="/app/wheels/src/flash-attn"
+  echo "[build] Preparando fonte FlashAttention (layer_norm) em ${SRC}"
+
+  # Sync/clone
+  if [ -d "$SRC/.git" ]; then
+    git -C "$SRC" fetch --all -p || true
+    git -C "$SRC" reset --hard origin/main || true
+    git -C "$SRC" clean -fdx || true
+  else
+    rm -rf "$SRC"
+    git clone --depth 1 https://github.com/Dao-AILab/flash-attention "$SRC"
+  fi
+
+  # Arquitetura CUDA alvo baseada na GPU ativa (uma só CC por build para wheel específica da máquina)
+  export TORCH_CUDA_ARCH_LIST="$(python - <<'PY'
+import torch
+try:
+    major, minor = torch.cuda.get_device_capability(0)
+    print(f"{major}.{minor}")
+except Exception:
+    # fallback genérico caso não haja GPU visível no build host
+    print("8.9")
+PY
+  )"
+  echo "[build] TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}"
+
+  # Compila apenas o submódulo csrc/layer_norm para wheel
+  pushd "$SRC/csrc/layer_norm" >/dev/null
+  export MAX_JOBS="${MAX_JOBS:-90}"
+  python -m pip wheel -v --no-build-isolation --no-deps . -w /app/wheels || true
+  popd >/dev/null
+
+  # Instala a wheel recém-criada
+  local W="$(ls -t /app/wheels/*flash*attn*layer*norm*-*.whl 2>/dev/null | head -n1 || true)"
+  if [ -z "${W}" ]; then
+    W="$(ls -t /app/wheels/*dropout*layer*norm*-*.whl 2>/dev/null | head -n1 || true)"
+  fi
+  if [ -z "${W}" ]; then
+    # fallback para qualquer wheel produzida no diretório
+    W="$(ls -t /app/wheels/*.whl 2>/dev/null | head -n1 || true)"
+  fi
+  if [ -n "${W}" ]; then
+    python -m pip install -v -U --no-deps "${W}" || true
+    echo "[build] FlashAttention LayerNorm instalado da wheel: ${W}"
+  else
+    echo "[build] Nenhuma wheel gerada; instalando do source (fallback)"
+    python -m pip install -v --no-build-isolation "$SRC/csrc/layer_norm" || true
+  fi
+}
+
 build_apex () {
   local SRC="/app/wheels/src/apex"
   echo "[build] Preparando fonte Apex em ${SRC}"
@@ -144,20 +260,21 @@ build_q8 () {
   fi
 }
 
-# ===== Pipeline genérico =====
+# ============================================================================
+#                          PIPELINES DE GARANTIA
+# ============================================================================
+
 ensure_pkg () {
   local PKG="$1"      # apex | q8_kernels
   local CHECK_FN="$2" # check_apex | check_q8
   local BUILD_FN="$3" # build_apex | build_q8
-
   echo "[flow] === ${PKG} ==="
   if ${CHECK_FN}; then
     echo "[flow] ${PKG}: já instalado (import OK)"
     return 0
   fi
-
   echo "[flow] ${PKG}: tentando wheel do Hub (${SELF_HF_REPO_ID})"
-  HF_OUT="$(install_from_hf "$PKG" || true)"
+  HF_OUT="$(install_from_hf_by_prefix "$PKG" || true)"
   if [ -n "${HF_OUT:-}" ]; then
     WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
     echo "[hub] Baixado: ${WHEEL_PATH}"
@@ -171,31 +288,71 @@ ensure_pkg () {
   else
     echo "[hub] Nenhuma wheel compatível encontrada para ${PKG}"
   fi
-
   echo "[flow] ${PKG}: compilando (fallback)"
   ${BUILD_FN}
   if ${CHECK_FN}; then
     echo "[flow] ${PKG}: sucesso após compilação"
     return 0
   fi
+  echo "[flow] ${PKG}: falhou após build; seguindo"
+  return 1
+}
 
-  echo "[flow] ${PKG}: falhou após build; registrando logs e seguindo"
+ensure_flashln () {
+  echo "[flow] === flash_attn_layer_norm ==="
+  if check_flashln; then
+    echo "[flow] FlashAttn LayerNorm: já instalado (import OK)"
+    return 0
+  fi
+  echo "[flow] FlashAttn LayerNorm: tentando wheel do Hub (${SELF_HF_REPO_ID})"
+  HF_OUT="$(install_flashln_from_hf || true)"
+  if [ -n "${HF_OUT:-}" ]; then
+    WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
+    echo "[hub] Baixado: ${WHEEL_PATH}"
+    python -m pip install -v -U --no-build-isolation "${WHEEL_PATH}" || true
+    if check_flashln; then
+      echo "[flow] FlashAttn LayerNorm: sucesso via Hub (${WHEEL_PATH})"
+      return 0
+    else
+      echo "[flow] FlashAttn LayerNorm: import falhou após wheel do Hub; compilando"
+    fi
+  else
+    echo "[hub] Nenhuma wheel compatível encontrada para FlashAttn LayerNorm"
+  fi
+  echo "[flow] FlashAttn LayerNorm: compilando (fallback)"
+  build_flashln
+  if check_flashln; then
+    echo "[flow] FlashAttn LayerNorm: sucesso após compilação"
+    return 0
+  fi
+  echo "[flow] FlashAttn LayerNorm: falhou após build; seguindo"
   return 1
 }
 
-# ===== Execução: Apex e Q8 =====
+# ============================================================================
+#                                 EXECUÇÃO
+# ============================================================================
+
+# 1) FlashAttention LayerNorm (novo cache em wheel no HF)
+ensure_flashln || true
+
+# 2) Apex (mantém pipeline existente)
 ensure_pkg "apex" check_apex build_apex || true
-#ensure_pkg "q8_kernels" check_q8 build_q8 || true
 
+# 3) Q8 kernels (opcional)
+# ensure_pkg "q8_kernels" check_q8 build_q8 || true
 
+# 4) Upload de wheels produzidas para o HF (cache cross-restarts)
 python - <<'PY'
 import os
 from huggingface_hub import HfApi, HfFolder
-repo=os.environ.get("SELF_HF_REPO_ID","euIaxs22/Aduc-sdr")
-token=os.getenv("HF_TOKEN") or HfFolder.get_token()
+
+repo = os.environ.get("SELF_HF_REPO_ID","euIaxs22/Aduc-sdr")
+token = os.getenv("HF_TOKEN") or HfFolder.get_token()
 if not token:
     raise SystemExit("HF_TOKEN ausente; upload desabilitado")
-api=HfApi(token=token)
+
+api = HfApi(token=token)
 api.upload_folder(
     folder_path="/app/wheels",
     repo_id=repo,
@@ -206,6 +363,5 @@ api.upload_folder(
 print("Upload concluído (wheels + licença).")
 PY
 
-
 chmod -R 777 /app/wheels || true
 echo "✅ Builder finalizado."