决赛现场脚本

tools/run_tn_dask_mpi_all.sh

@@ -11,25 +11,165 @@ NLAYERS="${NLAYERS:-20}"
 TORCH_THREADS="${TORCH_THREADS:-48}"
 SEARCH_REPEATS="${SEARCH_REPEATS:-2048}"
 SEARCH_TIME="${SEARCH_TIME:-300}"
-TN_TARGET_SIZE="${TN_TARGET_SIZE:-8589934592}"
+TN_TARGET_SIZE="${TN_TARGET_SIZE:-17179869184}"
 TN_TARGET_SLICES="${TN_TARGET_SLICES:-}"
 
 PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
 DTYPE="${DTYPE:-complex64}"
 TREE_DIR="${TREE_DIR:-trees/contest_tn}"
 DASK_ADDRESS="${DASK_ADDRESS:-tcp://10.20.1.103:8786}"
-MPIEXEC_FULL="${MPIEXEC_FULL:-mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2}"
+DASK_EXPECTED_WORKERS="${DASK_EXPECTED_WORKERS:-}"
+DASK_WAIT_FOR_WORKERS="${DASK_WAIT_FOR_WORKERS:-1}"
+DASK_WAIT_TIMEOUT="${DASK_WAIT_TIMEOUT:-600}"
+TN_DEBUG_TRIALS="${TN_DEBUG_TRIALS:-0}"
+MPIEXEC="${MPIEXEC:-mpirun}"
+MPIEXEC_FULL="${MPIEXEC_FULL:-}"
+MPI_HOSTS="${MPI_HOSTS:-}"
+MPI_HOSTFILE="${MPI_HOSTFILE:-${HOSTFILE:-}}"
+MPI_RANKS="${MPI_RANKS:-}"
+MPI_PE="${MPI_PE:-$TORCH_THREADS}"
+MPI_MAP_BY="${MPI_MAP_BY:-ppr:1:numa:PE=$MPI_PE}"
+MPI_BIND_TO="${MPI_BIND_TO:-core}"
+MPI_REPORT_BINDINGS="${MPI_REPORT_BINDINGS:-0}"
+MPI_EXPORT_ENV="${MPI_EXPORT_ENV:-1}"
+TN_CONTRACT_ENV_CHECK="${TN_CONTRACT_ENV_CHECK:-1}"
 SYNC_TREES="${SYNC_TREES:-1}"
 SYNC_HOSTS="${SYNC_HOSTS:-${WORKER_HOSTS:-}}"
 SSH_BIN="${SSH_BIN:-ssh}"
+DASK_CLUSTER_MANAGED="${DASK_CLUSTER_MANAGED:-0}"
 
 export TCM_ENABLE="${TCM_ENABLE:-1}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-$TORCH_THREADS}"
+export MKL_NUM_THREADS="${MKL_NUM_THREADS:-$TORCH_THREADS}"
+source "$ROOT_DIR/tools/qibotn_torch_mt_env.sh"
 
 tn_slice_args=(--tn-target-size "$TN_TARGET_SIZE")
 if [[ -n "$TN_TARGET_SLICES" ]]; then
   tn_slice_args+=(--tn-target-slices "$TN_TARGET_SLICES")
 fi
 
+cleanup_dask_cluster() {
+  local status=$?
+  if [[ "$DASK_CLUSTER_MANAGED" == "1" ]]; then
+    set +e
+    tools/manage_tn_dask_cluster.sh stop >/dev/null 2>&1 || true
+  fi
+  exit "$status"
+}
+
+trap cleanup_dask_cluster EXIT INT TERM HUP
+
+sum_host_slots() {
+  local hosts="$1"
+  local total=0
+  local item slots
+  IFS=',' read -r -a host_items <<< "$hosts"
+  for item in "${host_items[@]}"; do
+    if [[ "$item" == *:* ]]; then
+      slots="${item##*:}"
+    else
+      slots=1
+    fi
+    total=$((total + slots))
+  done
+  echo "$total"
+}
+
+count_hosts() {
+  local hosts="$1"
+  local count=0
+  local item
+  IFS=' ' read -r -a host_items <<< "$hosts"
+  for item in "${host_items[@]}"; do
+    [[ -n "$item" ]] && count=$((count + 1))
+  done
+  echo "$count"
+}
+
+wait_for_dask_workers() {
+  [[ "$DASK_WAIT_FOR_WORKERS" == "1" ]] || return 0
+  local expected="$DASK_EXPECTED_WORKERS"
+  if [[ -z "$expected" && -n "$WORKER_HOSTS" ]]; then
+    expected=$(( $(count_hosts "$WORKER_HOSTS") * NWORKERS ))
+  fi
+  if [[ -z "$expected" || "$expected" -le 0 ]]; then
+    return 0
+  fi
+
+  echo "Waiting for Dask workers: expected=$expected timeout=${DASK_WAIT_TIMEOUT}s"
+  "$PYTHON_BIN" - "$DASK_ADDRESS" "$expected" "$DASK_WAIT_TIMEOUT" <<'PY'
+import sys
+import time
+from distributed import Client
+
+address, expected, timeout = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
+deadline = time.time() + timeout
+client = Client(address)
+try:
+    while True:
+        info = client.scheduler_info(n_workers=-1)
+        workers = info.get("workers", {})
+        count = len(workers)
+        if count >= expected:
+            print(f"dask_workers_ready count={count} expected={expected}", flush=True)
+            break
+        if time.time() >= deadline:
+            print(
+                f"dask_workers_wait_timeout count={count} expected={expected}",
+                flush=True,
+            )
+            break
+        time.sleep(2)
+finally:
+    client.close()
+PY
+}
+
+append_mpi_env_args() {
+  [[ "$MPI_EXPORT_ENV" == "1" ]] || return 0
+  mpi_prefix+=(
+    -x "LD_PRELOAD=${LD_PRELOAD:-}"
+    -x "BLIS_NUM_THREADS=$BLIS_NUM_THREADS"
+    -x "OMP_NUM_THREADS=$OMP_NUM_THREADS"
+    -x "MKL_NUM_THREADS=$MKL_NUM_THREADS"
+    -x "OMP_PROC_BIND=$OMP_PROC_BIND"
+    -x "OMP_PLACES=$OMP_PLACES"
+  )
+}
+
+build_mpi_prefix() {
+  if [[ -n "$MPIEXEC_FULL" ]]; then
+    # shellcheck disable=SC2206
+    mpi_prefix=($MPIEXEC_FULL)
+    append_mpi_env_args
+    return
+  fi
+
+  local ranks="$MPI_RANKS"
+  if [[ -z "$ranks" && -n "$MPI_HOSTS" ]]; then
+    ranks="$(sum_host_slots "$MPI_HOSTS")"
+  fi
+  if [[ -z "$ranks" ]]; then
+    ranks=2
+  fi
+
+  mpi_prefix=(
+    "$MPIEXEC"
+    --map-by "$MPI_MAP_BY"
+    --bind-to "$MPI_BIND_TO"
+    -np "$ranks"
+  )
+  if [[ "$MPI_REPORT_BINDINGS" == "1" ]]; then
+    mpi_prefix+=(--report-bindings)
+  fi
+  append_mpi_env_args
+  if [[ -n "$MPI_HOSTS" ]]; then
+    mpi_prefix+=(-host "$MPI_HOSTS")
+  elif [[ -n "$MPI_HOSTFILE" ]]; then
+    mpi_prefix+=(-hostfile "$MPI_HOSTFILE")
+  fi
+}
+
 is_local_host() {
   local host="$1"
   [[ "$host" == "localhost" || "$host" == "127.0.0.1" ]] && return 0
@@ -62,25 +202,52 @@ sync_trees_to_hosts() {
 }
 
 tools/manage_tn_dask_cluster.sh start
+DASK_CLUSTER_MANAGED=1
+wait_for_dask_workers
 
 echo "Search with dask: $DASK_ADDRESS"
-"$PYTHON_BIN" -u tools/tn_contest_runner.py search \
-  --case "$CASE" \
-  --nqubits "$NQUBITS" \
-  --nlayers "$NLAYERS" \
-  --observables $OBSERVABLES \
-  --tree-dir "$TREE_DIR" \
-  --dask-address "$DASK_ADDRESS" \
-  --torch-threads "$TORCH_THREADS" \
-  --dtype "$DTYPE" \
-  --tn-search-repeats "$SEARCH_REPEATS" \
-  --tn-search-time "$SEARCH_TIME" \
+search_args=(
+  --case "$CASE"
+  --nqubits "$NQUBITS"
+  --nlayers "$NLAYERS"
+  --observables $OBSERVABLES
+  --tree-dir "$TREE_DIR"
+  --dask-address "$DASK_ADDRESS"
+  --torch-threads "$TORCH_THREADS"
+  --dtype "$DTYPE"
+  --tn-search-repeats "$SEARCH_REPEATS"
+  --tn-search-time "$SEARCH_TIME"
   "${tn_slice_args[@]}"
+)
+if [[ -n "$DASK_EXPECTED_WORKERS" ]]; then
+  search_args+=(--dask-expected-workers "$DASK_EXPECTED_WORKERS")
+fi
+if [[ "$TN_DEBUG_TRIALS" == "1" ]]; then
+  search_args+=(--tn-debug-trials)
+fi
+"$PYTHON_BIN" -u tools/tn_contest_runner.py search "${search_args[@]}"
 
 sync_trees_to_hosts
 
-echo "Contract with MPI: $MPIEXEC_FULL"
-read -r -a mpi_prefix <<< "$MPIEXEC_FULL"
+build_mpi_prefix
+echo "Contract with MPI: ${mpi_prefix[*]}"
+if [[ "$TN_CONTRACT_ENV_CHECK" == "1" ]]; then
+  "${mpi_prefix[@]}" "$PYTHON_BIN" -c "from mpi4py import MPI; import os; \
+import torch; \
+rank = MPI.COMM_WORLD.Get_rank(); \
+blis = []; \
+[blis.append(line.strip().split()[-1]) for line in open('/proc/self/maps') if 'libblis' in line and line.strip().split()[-1] not in blis]; \
+print('tn_contract_env ' + \
+      f'rank={rank} ' + \
+      f'LD_PRELOAD={os.environ.get(\"LD_PRELOAD\", \"\")} ' + \
+      f'BLIS_NUM_THREADS={os.environ.get(\"BLIS_NUM_THREADS\", \"\")} ' + \
+      f'OMP_NUM_THREADS={os.environ.get(\"OMP_NUM_THREADS\", \"\")} ' + \
+      f'MKL_NUM_THREADS={os.environ.get(\"MKL_NUM_THREADS\", \"\")} ' + \
+      f'OMP_PROC_BIND={os.environ.get(\"OMP_PROC_BIND\", \"\")} ' + \
+      f'OMP_PLACES={os.environ.get(\"OMP_PLACES\", \"\")} ' + \
+      f'torch_threads={torch.get_num_threads()} ' + \
+      f'blis={\";\".join(blis) if blis else \"missing\"}', flush=True)"
+fi
 "${mpi_prefix[@]}" "$PYTHON_BIN" -u tools/tn_contest_runner.py contract \
   --mpi \
   --case "$CASE" \