#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT

set -euo pipefail

# Wait for the NVIDIA driver to come up (poll for up to 60 seconds)
for _ in {1..60}; do
    nvidia-smi -L >/dev/null 2>&1 && break
    sleep 1
done

export CUDA_DEVICE_ORDER=PCI_BUS_ID

GB300_IDX=$(nvidia-smi --query-gpu=index,name --format=csv,noheader \
            | awk -F',' '/GB300/ {gsub(/ /,"",$1); print $1; exit}')

if [[ -z "${GB300_IDX:-}" ]]; then
    echo "GB300 not found" >&2
    exit 1
fi

# Persist for systemd EnvironmentFile + profile.d consumers
install -d -m 0755 /etc/cuda-gpu-select
cat >/etc/cuda-gpu-select/env <<EOT
CUDA_DEVICE_ORDER=PCI_BUS_ID
CUDA_VISIBLE_DEVICES=${GB300_IDX}
EOT
chmod 0644 /etc/cuda-gpu-select/env

# Atomically refresh /etc/environment so PAM logins see the values
tmp=$(mktemp)
grep -vE '^(CUDA_DEVICE_ORDER|CUDA_VISIBLE_DEVICES)=' /etc/environment 2>/dev/null > "$tmp" || true
cat /etc/cuda-gpu-select/env >> "$tmp"
install -m 0644 "$tmp" /etc/environment
rm -f "$tmp"