From 2a6de6794ddd4a35410a3f02d7a5c1936d24e123 Mon Sep 17 00:00:00 2001 From: yoshifuminakamura Date: Thu, 17 Sep 2026 17:27:30 +0900 Subject: [PATCH] Adjust SBD NCU profile defaults Signed-off-by: yoshifuminakamura --- programs/sbd/profile.sh | 141 +++++++++++++++++++++----- programs/sbd/run.sh | 2 +- scripts/tests/test_sbd_ncu_profile.sh | 44 +++++++- 3 files changed, 156 insertions(+), 31 deletions(-) diff --git a/programs/sbd/profile.sh b/programs/sbd/profile.sh index b936684..600304b 100644 --- a/programs/sbd/profile.sh +++ b/programs/sbd/profile.sh @@ -13,6 +13,43 @@ sbd_ncu_profile_mode() { printf '%s\n' "${BK_SBD_NCU_PROFILE_MODE:-discovery}" } +sbd_ncu_profile_timeout_seconds() { + local timeout_seconds="${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS:-0}" + + case "$timeout_seconds" in + ''|*[!0-9]*) + echo "SBD NCU profile timeout must be a non-negative integer: ${timeout_seconds}" >&2 + return 1 + ;; + esac + + printf '%s\n' "$timeout_seconds" +} + +sbd_strip_ncu_launch_count_args() { + local filtered_args=() + + while [ "$#" -gt 0 ]; do + case "$1" in + --launch-count) + shift + if [ "$#" -gt 0 ]; then + shift + fi + ;; + --launch-count=*) + shift + ;; + *) + filtered_args+=("$1") + shift + ;; + esac + done + + printf '%s\n' "${filtered_args[*]}" +} + sbd_supports_ncu_profile() { case "$1" in RIKYU|RC_DGXSP|RC_GH200) return 0 ;; @@ -24,6 +61,7 @@ sbd_configure_ncu_profile_from_run_env() { local system_name="$1" local profiler_tool local profiler_level + local default_profiler_level="detailed" if ! sbd_supports_ncu_profile "$system_name"; then return 0 @@ -44,8 +82,20 @@ sbd_configure_ncu_profile_from_run_env() { if [ -z "${BK_SBD_NCU_PROFILE:-}" ]; then export BK_SBD_NCU_PROFILE=true fi + if [ "$system_name" = "RIKYU" ]; then + default_profiler_level="single" + if [ -z "${BK_SBD_NCU_PLAN_TOP_K:-}" ]; then + export BK_SBD_NCU_PLAN_TOP_K=1 + fi + if [ -z "${BK_SBD_NCU_PLAN_LAUNCH_COUNT:-}" ]; then + export BK_SBD_NCU_PLAN_LAUNCH_COUNT=1 + fi + if [ -z "${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS:-}" ]; then + export BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS=1800 + fi + fi if [ -z "${BK_SBD_NCU_PROFILER_LEVEL:-}" ]; then - profiler_level=$(bk_resolve_profiler_level detailed SBD_PROFILER_LEVEL) + profiler_level=$(bk_resolve_profiler_level "$default_profiler_level" SBD_PROFILER_LEVEL) export BK_SBD_NCU_PROFILER_LEVEL="$profiler_level" fi } @@ -280,6 +330,8 @@ sbd_run_rank0_ncu_profile() { local report_file local profiler_status local archive_status + local profile_timeout_seconds + local profile_cmd=() results_dir=$(sbd_profile_results_dir) archive_path="${results_dir}/padata_${profile_slug}.tgz" @@ -298,39 +350,56 @@ sbd_run_rank0_ncu_profile() { ncu_level_arg_text=$(bk_profiler_ncu_level_args "$profiler_level") || return 1 read -r -a ncu_level_args <<< "$ncu_level_arg_text" + ncu_level_arg_text=$(sbd_strip_ncu_launch_count_args "${ncu_level_args[@]}") || return 1 + read -r -a ncu_level_args <<< "$ncu_level_arg_text" + profile_timeout_seconds=$(sbd_ncu_profile_timeout_seconds) || return 1 rm -rf "$raw_dir" "$stage_dir" mkdir -p "$rep_dir" "$stage_dir/raw" "$stage_dir/reports" echo "SBD NCU profile: profile='${profile_name}' kernel='${kernel_regex}' skip=${launch_skip} count=${launch_count}" >&2 echo "bk_profiler[ncu]: starting ${rep_name} level=${profiler_level} on rank 0" >&2 - set +e - mpirun -np "$n_ranks" bash -lc ' - rank=${OMPI_COMM_WORLD_RANK:-${PMIX_RANK:-${SLURM_PROCID:-0}}} - local_rank=${OMPI_COMM_WORLD_LOCAL_RANK:-${SLURM_LOCALID:-0}} - export CUDA_VISIBLE_DEVICES="${local_rank}" - profile_base="$1" - level_argc="$2" - shift 2 - level_args=() - i=0 - while [ "$i" -lt "$level_argc" ]; do - level_args+=("$1") - shift - i=$((i + 1)) - done - kernel_regex="$1" - launch_skip="$2" - launch_count="$3" - shift 3 - if [ "$rank" = 0 ]; then - exec ncu -o "$profile_base" --target-processes all "${level_args[@]}" \ - --kernel-name-base demangled --kernel-name "$kernel_regex" \ - --launch-skip "$launch_skip" --launch-count "$launch_count" ./diag "$@" + if [ "$profile_timeout_seconds" -gt 0 ]; then + echo "SBD NCU profile timeout: ${profile_timeout_seconds}s" >&2 + if ! command -v timeout >/dev/null 2>&1; then + echo "SBD NCU profile timeout requested but timeout command is not available." >&2 + return 1 fi - exec ./diag "$@" - ' bash "$profile_base" "${#ncu_level_args[@]}" "${ncu_level_args[@]}" \ - "$kernel_regex" "$launch_skip" "$launch_count" "$@" > "$profile_log" 2>&1 + fi + profile_cmd=( + mpirun -np "$n_ranks" bash -lc ' + rank=${OMPI_COMM_WORLD_RANK:-${PMIX_RANK:-${SLURM_PROCID:-0}}} + local_rank=${OMPI_COMM_WORLD_LOCAL_RANK:-${SLURM_LOCALID:-0}} + export CUDA_VISIBLE_DEVICES="${local_rank}" + profile_base="$1" + level_argc="$2" + shift 2 + level_args=() + i=0 + while [ "$i" -lt "$level_argc" ]; do + level_args+=("$1") + shift + i=$((i + 1)) + done + kernel_regex="$1" + launch_skip="$2" + launch_count="$3" + shift 3 + if [ "$rank" = 0 ]; then + exec ncu -o "$profile_base" --target-processes all "${level_args[@]}" \ + --kernel-name-base demangled --kernel-name "$kernel_regex" \ + --launch-skip "$launch_skip" --launch-count "$launch_count" ./diag "$@" + fi + exec ./diag "$@" + ' bash "$profile_base" "${#ncu_level_args[@]}" "${ncu_level_args[@]}" \ + "$kernel_regex" "$launch_skip" "$launch_count" "$@" + ) + set +e + if [ "$profile_timeout_seconds" -gt 0 ]; then + timeout --kill-after=60s "$profile_timeout_seconds" "${profile_cmd[@]}" > "$profile_log" 2>&1 + else + "${profile_cmd[@]}" > "$profile_log" 2>&1 + fi profiler_status=$? set -e @@ -455,6 +524,24 @@ PY fi } +sbd_run_optional_ncu_profiles() { + local system_name="$1" + local n_ranks="$2" + local profile_status + shift 2 + + if ! sbd_ncu_profile_enabled; then + return 0 + fi + + profile_status=0 + sbd_run_configured_ncu_profiles "$system_name" "$n_ranks" "$@" || profile_status=$? + if [ "$profile_status" -ne 0 ]; then + echo "SBD NCU profile acquisition failed with status ${profile_status}; continuing with benchmark result." >&2 + fi + return 0 +} + sbd_run_configured_ncu_profiles() { local system_name="$1" local n_ranks="$2" diff --git a/programs/sbd/run.sh b/programs/sbd/run.sh index 1b73428..336ff89 100755 --- a/programs/sbd/run.sh +++ b/programs/sbd/run.sh @@ -136,11 +136,11 @@ if ! awk -v actual="${energy}" -v reference="${reference_energy}" \ fi cp diag.log "${RESULTS_DIR}/" -sbd_run_configured_ncu_profiles "${system}" "${n_ranks}" "${diag_args[@]}" bk_emit_result --fom "${davidson_time}" --fom-unit s \ --fom-version "davidson_internal_s" --exp "${experiment}" \ --nodes "${nodes}" --numproc-node "${numproc_node}" \ --nthreads "${nthreads}" >> "${RESULTS_DIR}/result" +sbd_run_optional_ncu_profiles "${system}" "${n_ranks}" "${diag_args[@]}" if [[ -n "${mult_time}" ]]; then bk_emit_section mult "${mult_time}" "" "${SBD_MULT_SECTION_ARTIFACTS:-}" >> "${RESULTS_DIR}/result" fi diff --git a/scripts/tests/test_sbd_ncu_profile.sh b/scripts/tests/test_sbd_ncu_profile.sh index 158e46b..6353d8c 100644 --- a/scripts/tests/test_sbd_ncu_profile.sh +++ b/scripts/tests/test_sbd_ncu_profile.sh @@ -103,6 +103,8 @@ fi output="" kernel="" +launch_count="" +launch_count_seen=0 while [ $# -gt 0 ]; do case "$1" in -o|--output) @@ -113,7 +115,16 @@ while [ $# -gt 0 ]; do shift kernel="$1" ;; - --kernel-name-base|--launch-skip|--launch-count|--target-processes|--set) + --launch-count) + launch_count_seen=$((launch_count_seen + 1)) + shift + launch_count="$1" + ;; + --launch-count=*) + launch_count_seen=$((launch_count_seen + 1)) + launch_count="${1#--launch-count=}" + ;; + --kernel-name-base|--launch-skip|--target-processes|--set) shift ;; --nvtx) @@ -127,8 +138,11 @@ while [ $# -gt 0 ]; do done test -n "$output" +if [ "${FAKE_NCU_FAIL:-}" = "1" ]; then + exit 7 +fi mkdir -p "$(dirname "$output")" -printf '%s\n' "$kernel" >> "$FAKE_NCU_LOG" +printf '%s\tlaunch_count=%s\tlaunch_count_seen=%s\n' "$kernel" "$launch_count" "$launch_count_seen" >> "$FAKE_NCU_LOG" printf 'fake ncu report\n' > "${output}.ncu-rep" EOF @@ -154,8 +168,30 @@ export PATH="${FAKE_BIN}:${PATH}" test "${BK_SBD_NCU_PROFILE:-}" = "false" ) +( + unset BK_PROFILER BK_PROFILER_LEVEL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL + unset BK_SBD_NCU_PLAN_TOP_K BK_SBD_NCU_PLAN_LAUNCH_COUNT BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS + unset SBD_PROFILER_TOOL + sbd_configure_ncu_profile_from_run_env RIKYU + test "${BK_SBD_NCU_PROFILE}" = "true" + test "${BK_SBD_NCU_PROFILER_LEVEL}" = "single" + test "${BK_SBD_NCU_PLAN_TOP_K}" = "1" + test "${BK_SBD_NCU_PLAN_LAUNCH_COUNT}" = "1" + case "${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS}" in + ''|0|*[!0-9]*) exit 1 ;; + esac +) + +test "$(sbd_strip_ncu_launch_count_args --set basic --launch-count 1 --nvtx --launch-count=3)" = "--set basic --nvtx" + +( + export BK_SBD_NCU_PROFILE=true + sbd_run_configured_ncu_profiles() { return 7; } + sbd_run_optional_ncu_profiles RIKYU 4 +) + pushd "${RUN_DIR}" >/dev/null -unset BK_PROFILER SBD_PROFILER_TOOL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL +unset BK_PROFILER SBD_PROFILER_TOOL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL BK_SBD_NCU_PLAN_LAUNCH_COUNT export BK_SBD_NCU_PROFILE_MODE=discovery export BK_SBD_NCU_PLAN_TOP_K=5 export BK_PROFILER_LEVEL=single @@ -174,6 +210,7 @@ test -f "${RESULTS_DIR}/sbd_ncu_plan.json" jq -e ' (.profiles | length) == 5 and + ([.profiles[].launch_count] | unique) == [1] and ([.profiles[].kernel_match.pattern] | unique | length) == 5 and .profiles[1].kernel_match.pattern == "regex:.*sbd::MultUnified.*" and .profiles[4].kernel_match.pattern == "regex:.*sbd::MultUnified.*" @@ -208,5 +245,6 @@ for archive in "${profile_archives[@]}"; do done grep -Fq 'regex:.*sbd::MultUnified.*' "$FAKE_NCU_LOG" +awk -F '\t' '$2 != "launch_count=1" || $3 != "launch_count_seen=1" { bad = 1 } END { exit bad }' "$FAKE_NCU_LOG" echo "SBD NCU profile test passed"