#!/usr/bin/env bash
# TE 2.14.0 wants cublasLtGroupedMatrixLayoutInit_internal, absent even from the
# newest nvidia-cublas 13.1.1.3. The known-good b300 stack (jihuny) runs TE
# 2.14.1, so match it. --no-deps keeps the torch 2.12.0+cu130 pin intact, which
# a plain install has already disturbed twice.
set -uo pipefail
RT=/KRAFTON/WORKSPACE/wbl-workspace/posttraining-2606/.b300-runtime
PY=$RT/bin/python; LOG=$RT/restore.log
log(){ echo "[$(date '+%F %H:%M:%S')] FIX3 $*" | tee -a "$LOG"; }
log "upgrading transformer-engine 2.14.0 -> 2.14.1"
$PY -m pip install --no-cache-dir --no-deps --upgrade \
  transformer_engine==2.14.1 transformer_engine_cu13==2.14.1 transformer_engine_torch==2.14.1 >>"$LOG" 2>&1
log "te rc=$?"
log "=== verification ==="
$PY - <<'PY' 2>&1 | tee -a "$LOG"
import importlib
for m in ["torch","transformer_engine","megatron.core","megatron.bridge","transformers","mamba_ssm","causal_conv1d"]:
    try:
        mod=importlib.import_module(m); print(f"  OK   {m:20s} {getattr(mod,'__version__','')}")
    except Exception as e: print(f"  MISS {m:20s} {type(e).__name__}: {str(e)[:70]}")
PY
log "FIX3 DONE"