|
| 1 | +#!/usr/bin/env bash |
| 2 | +# Hygon/DTK environment setup for TransformerEngine-FL plugin QA. |
| 3 | +set -euo pipefail |
| 4 | + |
| 5 | +WORKSPACE="${GITHUB_WORKSPACE:-$(pwd)}" |
| 6 | + |
| 7 | +echo "===== Load Hygon/DTK runtime environment =====" |
| 8 | +source "$WORKSPACE/tests/plugin/backend/hygon/set_env.sh" |
| 9 | + |
| 10 | +# Hygon CI is a reference-backend baseline. Force the selection policy here so |
| 11 | +# inherited shell state cannot silently fall back to FlagOS. |
| 12 | +export TE_FL_SKIP_CUDA=1 |
| 13 | +export TE_FL_PREFER=reference |
| 14 | +export NVTE_FRAMEWORK=pytorch |
| 15 | +export NVTE_FLASH_ATTN=0 |
| 16 | +export NVTE_FUSED_ATTN=0 |
| 17 | +export NVTE_UNFUSED_ATTN=1 |
| 18 | +export NVTE_UnfusedDPA_Emulate_FP8=1 |
| 19 | + |
| 20 | +echo "===== Verify Hygon device visibility =====" |
| 21 | +if [ "${HYGON_REQUIRE_DEVICE:-1}" = "1" ] && ! command -v hy-smi >/dev/null 2>&1; then |
| 22 | + echo "ERROR: hy-smi is unavailable in the Hygon CI image" >&2 |
| 23 | + exit 1 |
| 24 | +elif command -v hy-smi >/dev/null 2>&1; then |
| 25 | + hy-smi |
| 26 | +else |
| 27 | + echo "WARNING: hy-smi is unavailable; device verification is disabled" |
| 28 | +fi |
| 29 | + |
| 30 | +echo "===== Verify Python runtime =====" |
| 31 | +"$PYTHON_BIN" - <<'PY' |
| 32 | +import os |
| 33 | +import sys |
| 34 | +
|
| 35 | +print("python:", sys.executable) |
| 36 | +print("version:", sys.version) |
| 37 | +
|
| 38 | +try: |
| 39 | + import torch |
| 40 | +except ModuleNotFoundError as exc: |
| 41 | + raise SystemExit(f"PyTorch is required in the Hygon CI image: {exc}") from exc |
| 42 | +
|
| 43 | +print("torch:", torch.__version__) |
| 44 | +
|
| 45 | +if os.environ.get("HYGON_REQUIRE_DEVICE", "1") == "1": |
| 46 | + if not torch.cuda.is_available(): |
| 47 | + raise SystemExit("Hygon DCU is not visible through torch.cuda") |
| 48 | +
|
| 49 | + device_count = torch.cuda.device_count() |
| 50 | + if device_count < 1: |
| 51 | + raise SystemExit("torch.cuda reports zero Hygon devices") |
| 52 | +
|
| 53 | + device = torch.device("cuda") |
| 54 | + lhs = torch.ones((2, 2), device=device) |
| 55 | + rhs = torch.full((2, 2), 2.0, device=device) |
| 56 | + result = lhs @ rhs |
| 57 | + if not torch.allclose(result.cpu(), torch.full((2, 2), 4.0)): |
| 58 | + raise SystemExit("Hygon DCU matrix-multiplication smoke test failed") |
| 59 | +
|
| 60 | + print("cuda_device_count:", device_count) |
| 61 | + print("cuda_device_name:", torch.cuda.get_device_name(0)) |
| 62 | + print("matmul_smoke: passed") |
| 63 | +PY |
| 64 | + |
| 65 | +echo "===== Verify reference backend selection =====" |
| 66 | +"$PYTHON_BIN" - <<'PY' |
| 67 | +from transformer_engine.plugin.core import get_manager |
| 68 | +
|
| 69 | +manager = get_manager() |
| 70 | +selected_impl = manager.get_selected_impl_id("generic_gemm") |
| 71 | +if selected_impl != "reference.torch": |
| 72 | + raise SystemExit( |
| 73 | + f"Expected generic_gemm to use reference.torch, selected {selected_impl!r}" |
| 74 | + ) |
| 75 | +print("generic_gemm_impl:", selected_impl) |
| 76 | +PY |
| 77 | + |
| 78 | +echo "===== Install Hygon QA dependencies =====" |
| 79 | +if [ "${HYGON_SKIP_DEP_INSTALL:-0}" = "1" ]; then |
| 80 | + echo "Skipping Python dependency installation because HYGON_SKIP_DEP_INSTALL=1" |
| 81 | +else |
| 82 | + missing_modules=() |
| 83 | + for module_name in pytest expecttest coverage pytest_cov; do |
| 84 | + if ! "$PYTHON_BIN" -c "import importlib; importlib.import_module('$module_name')" >/dev/null 2>&1; then |
| 85 | + missing_modules+=("$module_name") |
| 86 | + fi |
| 87 | + done |
| 88 | + |
| 89 | + if [ "${#missing_modules[@]}" -gt 0 ]; then |
| 90 | + echo "Missing Hygon QA modules: ${missing_modules[*]}" |
| 91 | + "$PYTHON_BIN" -m pip install pytest==8.2.1 expecttest coverage pytest-cov |
| 92 | + else |
| 93 | + echo "Hygon QA dependencies are already available in the image" |
| 94 | + fi |
| 95 | + |
| 96 | + # ONNX dependencies are installed only by the ONNX test group. |
| 97 | +fi |
| 98 | + |
| 99 | +"$PYTHON_BIN" -c "import coverage, pytest_cov; print('coverage dependencies: ready')" |
| 100 | + |
| 101 | +if [ "${HYGON_INSTALL_TE:-0}" = "1" ]; then |
| 102 | + echo "===== Install TransformerEngine-FL Python layer =====" |
| 103 | + cd "$WORKSPACE" |
| 104 | + TE_FL_SKIP_CUDA=1 "$PYTHON_BIN" setup.py install |
| 105 | +else |
| 106 | + echo "Skipping TransformerEngine-FL install; tests run from source via PYTHONPATH" |
| 107 | +fi |
| 108 | + |
| 109 | +if [ -n "${GITHUB_ENV:-}" ]; then |
| 110 | + { |
| 111 | + echo "PATH=$PATH" |
| 112 | + echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}" |
| 113 | + echo "PYTHONPATH=$WORKSPACE${PYTHONPATH:+:$PYTHONPATH}" |
| 114 | + echo "TE_PATH=$WORKSPACE" |
| 115 | + echo "XML_LOG_DIR=$WORKSPACE/logs" |
| 116 | + echo "PLATFORM=$PLATFORM" |
| 117 | + echo "TE_FL_SKIP_CUDA=$TE_FL_SKIP_CUDA" |
| 118 | + echo "TE_FL_PREFER=$TE_FL_PREFER" |
| 119 | + echo "NVTE_FRAMEWORK=$NVTE_FRAMEWORK" |
| 120 | + echo "PYTHON_BIN=$PYTHON_BIN" |
| 121 | + echo "NVTE_FLASH_ATTN=$NVTE_FLASH_ATTN" |
| 122 | + echo "NVTE_FUSED_ATTN=$NVTE_FUSED_ATTN" |
| 123 | + echo "NVTE_UNFUSED_ATTN=$NVTE_UNFUSED_ATTN" |
| 124 | + echo "NVTE_UnfusedDPA_Emulate_FP8=$NVTE_UnfusedDPA_Emulate_FP8" |
| 125 | + echo "HYGON_REQUIRE_DEVICE=${HYGON_REQUIRE_DEVICE:-1}" |
| 126 | + } >> "$GITHUB_ENV" |
| 127 | +fi |
| 128 | + |
| 129 | +echo "===== Hygon environment setup complete =====" |
0 commit comments