[Slurm] Collect more data for Slurm nodes

src/cloudai/systems/slurm/slurm-metadata.sh

@@ -1,35 +1,141 @@
-. /etc/os-release
-hpcx_version=${HPCX_DIR##*/}
+#!/usr/bin/env bash
 
-while IFS=: read -r value; do
-    echo "$value"
-done <<EOF
-user = "${USER:-$(whoami)}"
+export LC_ALL=C
+. "${CLOUDAI_OS_RELEASE_FILE:-/etc/os-release}" 2>/dev/null || true
+
+inventory() {
+    awk 'NF { count[$0]++; if (count[$0] == 1) order[++n] = $0 }
+         END { for (i = 1; i <= n; i++) printf "%s%s x%d", (i > 1 ? ", " : ""), order[i], count[order[i]] }'
+}
+
+first_or_null() {
+    awk 'NF { print; found = 1; exit } END { if (!found) print "null" }'
+}
+
+ofed_version() {
+    command -v ofed_info >/dev/null || { echo null; return; }
+    ofed_info -s 2>/dev/null | sed 's/:$//' | first_or_null
+}
+
+libfabric_version() {
+    command -v fi_info >/dev/null || { echo null; return; }
+    fi_info --version 2>/dev/null |
+        awk '{ for (i = 1; i <= NF; i++) if ($i ~ /^[0-9]+([.][0-9]+)+/) { print $i; exit } }' |
+        first_or_null
+}
+
+cuda_toolkit_version() {
+    local version
+    version="$(nvcc --version 2>/dev/null | sed -n 's/.*release \([0-9.]*\).*/\1/p' | head -n1)"
+    [ -n "$version" ] || version="$(sed -n 's/.*"version"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' \
+        "${CLOUDAI_CUDA_ROOT:-/usr/local/cuda}/version.json" 2>/dev/null | head -n1)"
+    [ -n "$version" ] || version="$(sed -n 's/[^0-9]*\([0-9][0-9.]*\).*/\1/p' \
+        "${CLOUDAI_CUDA_ROOT:-/usr/local/cuda}/version.txt" 2>/dev/null | head -n1)"
+    printf '%s' "${version:-null}"
+}
+
+doca_host_version() {
+    local doca_root="${CLOUDAI_DOCA_ROOT:-/opt/mellanox/doca}"
+    local doca_info version
 
+    doca_info="$(command -v doca-info 2>/dev/null)"
+    doca_info="${doca_info:-$doca_root/tools/doca-info}"
+    if [ -x "$doca_info" ]; then
+        version="$("$doca_info" 2>/dev/null | awk '
+            /^DOCA:/ { found = 1; next }
+            found && /^- / { for (i = 1; i <= NF; i++) if ($i ~ /^[0-9]+\.[0-9]+/) { print $i; exit } }
+        ')"
+    fi
+    if [ -z "$version" ] && command -v dpkg-query >/dev/null 2>&1; then
+        version="$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}\n' 'doca-*' 2>/dev/null |
+            awk '$1 ~ /^ii/ { print $2; exit }')"
+    fi
+    if [ -z "$version" ] && command -v rpm >/dev/null 2>&1; then
+        version="$(rpm -qa --qf '%{VERSION}-%{RELEASE}\n' 'doca-*' 2>/dev/null | head -n1)"
+    fi
+    if [ -z "$version" ]; then
+        version="$(sed -n 's/.*"version"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' \
+            "$doca_root/version.json" 2>/dev/null | head -n1)"
+    fi
+    [ -n "$version" ] || version="$(awk 'NF { print; exit }' "$doca_root/version.txt" 2>/dev/null)"
+    printf '%s' "${version:-null}"
+}
+
+system_metadata() {
+    local gpu_names gpu_model gpu_count gpu_inventory
+    local cpu_model cpu_arch cpu_vendor
+
+    gpu_names="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null)"
+    gpu_model="$(printf '%s\n' "$gpu_names" | awk 'NF { print; exit }')"
+    gpu_count="$(printf '%s\n' "$gpu_names" | awk 'NF { count++ } END { print count + 0 }')"
+    gpu_inventory="$(printf '%s\n' "$gpu_names" | inventory)"
+    cpu_model="$(lscpu 2>/dev/null | sed -n 's/^Model name:[[:space:]]*//p' | head -n1)"
+    cpu_arch="$(lscpu 2>/dev/null | sed -n 's/^Architecture:[[:space:]]*//p' | head -n1)"
+    cpu_vendor="$(lscpu 2>/dev/null | sed -n 's/^Vendor ID:[[:space:]]*//p' | head -n1)"
+
+    cat <<EOF
 [system]
 os_type = "${ID:-null}"
-os_version = "${VERSION:-null}"
-linux_kernel_version = "$(uname -r)"
-gpu_arch_type = "$(nvidia-smi -q 2>/dev/null | grep "Product Name" | head -n1 | awk -F': ' '{print $2}' || echo null)"
-cpu_model_name = "$(lscpu 2>/dev/null | grep "Model name:" | sed 's/Model name:[[:space:]]*//' || echo null)"
-cpu_arch_type = "$(lscpu 2>/dev/null | grep "Architecture:" | sed 's/Architecture:[[:space:]]*//' || echo null)"
+os_version = "${VERSION:-${VERSION_ID:-null}}"
+linux_kernel_version = "$(uname -r 2>/dev/null || echo null)"
+gpu_arch_type = "${gpu_model:-null}"
+gpu_count = $gpu_count
+gpu_inventory = "${gpu_inventory:-null}"
+cpu_model_name = "${cpu_model:-null}"
+cpu_arch_type = "${cpu_arch:-null}"
+cpu_vendor = "${cpu_vendor:-null}"
+EOF
+}
+
+physical_network_metadata() {
+    local nic_lines nic_models nics nic_count nic_inventory firmware
+    local fw_file device
+
+    nic_lines="$(lspci -Dnn 2>/dev/null | awk '
+        tolower($0) ~ /ethernet controller|infiniband controller/ &&
+        (tolower($0) ~ /mellanox|nvidia/ || tolower($0) ~ /\[15b3:/)
+    ')"
+    nic_models="$(printf '%s\n' "$nic_lines" | sed -E 's/^[^ ]+ [^:]+: //')"
+    nics="$(printf '%s\n' "$nic_models" | awk 'NF { print; exit }')"
+    nic_count="$(printf '%s\n' "$nic_lines" | awk 'NF { count++ } END { print count + 0 }')"
+    nic_inventory="$(printf '%s\n' "$nic_models" | inventory)"
+
+    for fw_file in "${CLOUDAI_INFINIBAND_SYSFS:-/sys/class/infiniband}"/*/fw_ver; do
+        [ -r "$fw_file" ] || continue
+        device="${fw_file%/fw_ver}"
+        firmware="${firmware:+$firmware, }${device##*/}=$(head -n1 "$fw_file")"
+    done
+
+    cat <<EOF
+nics = "${nics:-null}"
+nic_count = $nic_count
+nic_inventory = "${nic_inventory:-null}"
+hca_firmware_versions = "${firmware:-null}"
+switch_type = "${SWITCH:-null}"
+network_name = "${NETWORK:-null}"
+mofed_version = "$(ofed_version)"
+doca_host_version = "$(doca_host_version)"
+EOF
+}
+
+runtime_metadata() {
+    local driver_version hpcx_version=${HPCX_DIR##*/}
+    driver_version="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1)"
+
+    cat <<EOF
+user = "${USER:-$(whoami)}"
 
 [mpi]
-mpi_type = "$(mpirun --version 2>/dev/null | grep -i "open mpi" -q && echo openmpi || echo null)"
-mpi_version = "$(mpirun --version 2>/dev/null | grep -oP "(?<=\(Open MPI\) )[^\s]+$" || echo null)"
+mpi_type = "$(mpirun --version 2>/dev/null | grep -qi 'open mpi' && echo openmpi || echo null)"
+mpi_version = "$(mpirun --version 2>/dev/null | awk 'tolower($0) ~ /open mpi/ { print $NF; exit }' | first_or_null)"
 hpcx_version = "${hpcx_version:-null}"
 
 [cuda]
 cuda_build_version = "${CUDA_BUILD_VERSION:-null}"
-cuda_runtime_version = "$(nvidia-smi 2>/dev/null | grep -oP "(?<=CUDA Version: )[\d\.]+" || echo null)"
-cuda_driver_version = "$(nvidia-smi 2>/dev/null | grep -oP "(?<=Driver Version: )[\d\.]+" || echo null)"
-
-[network]
-nics = "$(lspci 2>/dev/null | grep -E "Ethernet controller|Infiniband controller" | awk -F": " '/Mellanox/ {found=1;print $NF;exit 0} END {if (!found) exit 1}' || echo null)"
-switch_type = "${SWITCH:-null}"
-network_name = "${NETWORK:-null}"
-mofed_version = "$(command -v ofed_info >/dev/null && ofed_info -s 2>/dev/null | sed 's/:$//' || echo null)"
-libfabric_version = "$(command -v fi_info >/dev/null && fi_info --version 2>/dev/null | grep "Libfabric" | awk '{print $2}' || echo null)"
+cuda_runtime_version = "$(nvidia-smi 2>/dev/null | sed -n 's/.*CUDA Version: \([0-9.]*\).*/\1/p' | first_or_null)"
+cuda_driver_version = "${driver_version:-null}"
+nvidia_driver_version = "${driver_version:-null}"
+cuda_toolkit_version = "$(cuda_toolkit_version)"
 
 [nccl]
 version = "${NCCL_VERSION:-null}"
@@ -42,4 +148,27 @@ num_nodes = "${SLURM_NNODES:-null}"
 ntasks_per_node = "${SLURM_NTASKS_PER_NODE:-null}"
 ntasks = "${SLURM_NTASKS:-null}"
 job_id = "${SLURM_JOBID:-null}"
+
+[network]
+libfabric_version = "$(libfabric_version)"
 EOF
+}
+
+case "${1:-all}" in
+    runtime)
+        runtime_metadata
+        ;;
+    host)
+        echo
+        physical_network_metadata
+        echo
+        system_metadata
+        ;;
+    *)
+        runtime_metadata
+        echo
+        physical_network_metadata
+        echo
+        system_metadata
+        ;;
+esac

src/cloudai/systems/slurm/slurm_command_gen_strategy.py

@@ -310,20 +310,50 @@ def _ranks_mapping_cmd(self) -> str:
     def _metadata_cmd(self) -> str:
         (self.test_run.output_path.absolute() / "metadata").mkdir(parents=True, exist_ok=True)
         num_nodes, _ = self.get_cached_nodes_spec()
-        metadata_script_path = "/cloudai_install"
-        if not self.image_path():
-            metadata_script_path = str(self.system.install_path.absolute())
-        return " ".join(
+        metadata_script_path = self.system.install_path.absolute()
+
+        metadata_output = self.test_run.output_path.absolute() / "metadata" / "node-%N.toml"
+        metadata_error = self.test_run.output_path.absolute() / "metadata" / "nodes.err"
+
+        metadata_srun_prefix = ["srun", "--export=ALL", "--mpi=none"]
+        if not self.nodelist_in_use:
+            metadata_srun_prefix.append(f"-N{num_nodes}")
+        if self.system.extra_srun_args:
+            metadata_srun_prefix.append(self.system.extra_srun_args)
+        if self.test_run.extra_srun_args:
+            metadata_srun_prefix.append(self.test_run.extra_srun_args)
+
+        metadata_srun_prefix = [
+            *metadata_srun_prefix,
+            f"--ntasks={num_nodes}",
+            "--ntasks-per-node=1",
+            f"--output={metadata_output}",
+            f"--error={metadata_error}",
+        ]
+
+        image_path = self.image_path()
+        if not image_path:
+            return " ".join([*metadata_srun_prefix, "bash", f"{metadata_script_path}/slurm-metadata.sh"])
+
+        runtime_command_parts = [*metadata_srun_prefix, f"--container-image={image_path}"]
+        mounts = self.container_mounts()
+        if mounts:
+            runtime_command_parts.append(f"--container-mounts={','.join(mounts)}")
+        if not self.system.container_mount_home:
+            runtime_command_parts.append("--no-container-mount-home")
+        runtime_command_parts.extend(["bash", f"{self.CONTAINER_MOUNT_INSTALL}/slurm-metadata.sh", "runtime"])
+        runtime_command = " ".join(runtime_command_parts)
+
+        host_command = " ".join(
             [
-                *self.gen_srun_prefix(),
-                f"--ntasks={num_nodes}",
-                "--ntasks-per-node=1",
-                f"--output={self.test_run.output_path.absolute() / 'metadata' / 'node-%N.toml'}",
-                f"--error={self.test_run.output_path.absolute() / 'metadata' / 'nodes.err'}",
+                *metadata_srun_prefix,
+                "--open-mode=append",
                 "bash",
                 f"{metadata_script_path}/slurm-metadata.sh",
+                "host",
             ]
         )
+        return "\n".join([runtime_command, host_command])
 
     def _enable_vboost_cmd(self) -> str:
         num_nodes, _ = self.system.get_nodes_by_spec(self.test_run.nnodes, self.test_run.nodes)

src/cloudai/systems/slurm/slurm_metadata.py

@@ -94,8 +94,11 @@ class MetadataSystem(BaseModel):
     os_version: str
     linux_kernel_version: str
     gpu_arch_type: str
+    gpu_count: int = 0
+    gpu_inventory: str = "null"
     cpu_model_name: str
     cpu_arch_type: str
+    cpu_vendor: str = "null"
 
 
 class MetadataMPI(BaseModel):
@@ -112,15 +115,21 @@ class MetadataCUDA(BaseModel):
     cuda_build_version: str
     cuda_runtime_version: str
     cuda_driver_version: str
+    nvidia_driver_version: str = "null"
+    cuda_toolkit_version: str = "null"
 
 
 class MetadataNetwork(BaseModel):
     """Represents the network metadata."""
 
     nics: str
+    nic_count: int = 0
+    nic_inventory: str = "null"
+    hca_firmware_versions: str = "null"
     switch_type: str
     network_name: str
     mofed_version: str
+    doca_host_version: str = "null"
     libfabric_version: str
 
 

tests/ref_data/ai-dynamo.sbatch

@@ -12,7 +12,8 @@ export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head
 
 srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__INSTALL_DIR__:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__/huggingface:/cloudai_install/huggingface,/tmp:/tmp --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__INSTALL_DIR__:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__/huggingface:/cloudai_install/huggingface,/tmp:/tmp --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__INSTALL_DIR__:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__/huggingface:/cloudai_install/huggingface,/tmp:/tmp bash /cloudai_install/slurm-metadata.sh runtime
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --open-mode=append bash __INSTALL_DIR__/slurm-metadata.sh host
 
 # Start DCGM exporter on each node.
 echo "Starting DCGM exporter..."

tests/ref_data/ddlb.sbatch

@@ -12,6 +12,7 @@ export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head
 
 srun --export=ALL --mpi=pmix -N1 --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL --mpi=pmix -N1 --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=none -N1 --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash /cloudai_install/slurm-metadata.sh runtime
+srun --export=ALL --mpi=none -N1 --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --open-mode=append bash __INSTALL_DIR__/slurm-metadata.sh host
 
 srun --export=ALL --mpi=pmix -N1 --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; python ddlb/cli/benchmark.py --primitive tp_columnwise -m 1024 -n 128 -k 1024 --dtype float16 --num-iterations 50 --num-warmups 5 --impl pytorch;backend=nccl;order=AG_before"

tests/ref_data/deepep-benchmark.sbatch

@@ -25,6 +25,7 @@ export MASTER_PORT=29500
 
 srun --export=ALL --mpi=pmix -N2 --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL --mpi=pmix -N2 --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash /cloudai_install/slurm-metadata.sh runtime
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --open-mode=append bash __INSTALL_DIR__/slurm-metadata.sh host
 
 srun --export=ALL --mpi=pmix --container-image=docker/image:url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --nodes=2 --ntasks=2 --ntasks-per-node=1 bash -c "source __OUTPUT_DIR__/output/env_vars.sh; torchrun --nnodes=2 --nproc_per_node=1 --rdzv_id=\${SLURM_JOB_ID:-0} --rdzv_backend=c10d --rdzv_endpoint=\${MASTER_ADDR}:\${MASTER_PORT} /workspace/DeepEP/tests/legacy/test_internode.py --num-processes 8 --num-tokens 4096 --hidden 7168 --num-topk 8 --pressure-test-mode 0 --num-experts 256"

tests/ref_data/fio.sbatch

@@ -12,6 +12,7 @@ export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head
 
 srun --export=ALL -N2 --container-image=openeuler/fio:3.42-oe2403sp3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL -N2 --container-image=openeuler/fio:3.42-oe2403sp3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --container-image=openeuler/fio:3.42-oe2403sp3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash /cloudai_install/slurm-metadata.sh runtime
+srun --export=ALL --mpi=none -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err --open-mode=append bash __INSTALL_DIR__/slurm-metadata.sh host
 
 srun --export=ALL -N2 --container-image=openeuler/fio:3.42-oe2403sp3 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks-per-node=2 --ntasks=4 bash -c "source __OUTPUT_DIR__/output/env_vars.sh; fio --name=fio-smoke --filename=/tmp/cloudai-fio-test --rw=randwrite --bs=128k --iodepth=1 --numjobs=1 --group_reporting --thread"