Adjust SBD NCU profile defaults

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"

programs/sbd/run.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
 set -euo pipefail
 
@@ -136,11 +136,11 @@
 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

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<double,[[:space:]]*\\(int\\)0,[[:space:]]*\\(int\\)1>.*" and
   .profiles[4].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)1,[[:space:]]*\\(int\\)0>.*"
@@ -208,5 +245,6 @@ for archive in "${profile_archives[@]}"; do
 done
 
 grep -Fq 'regex:.*sbd::MultUnified<double,[[:space:]]*\(int\)0,[[:space:]]*\(int\)1>.*' "$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"