@@ -34,9 +34,7 @@ echo ""
3434
3535# Use $(pwd) not BASH_SOURCE — sbatch copies the script to /var/spool/slurmd/
3636# at runtime, but the CWD remains the submit-time directory (amd_utils/).
37- if [[ " $ENGINE " == " atom-disagg" ]]; then
38- MODELS_YAML=" $( pwd) /models_atom.yaml"
39- elif [[ " $ENGINE " == " tilert" ]]; then
37+ if [[ " $ENGINE " == " tilert" ]]; then
4038 MODELS_YAML=" $( pwd) /models_tilert.yaml"
4139else
4240 MODELS_YAML=" $( pwd) /models.yaml"
@@ -543,26 +541,7 @@ DOCKER_ENV_COMMON=(
543541)
544542
545543# Engine-specific env vars
546- if [[ " $ENGINE " == " atom-disagg" ]]; then
547- check_env_vars \
548- PREFILL_PORT DECODE_PORT HANDSHAKE_PORT MEM_FRAC_STATIC KV_CACHE_DTYPE \
549- BLOCK_SIZE MAX_NUM_SEQS
550- DOCKER_ENV_ENGINE=(
551- -e ATOM_WS_PATH=${WS_PATH}
552- -e PREFILL_PORT=${PREFILL_PORT}
553- -e DECODE_PORT=${DECODE_PORT}
554- -e ROUTER_PORT=${ROUTER_PORT}
555- -e HANDSHAKE_PORT=${HANDSHAKE_PORT}
556- -e MEM_FRAC_STATIC=${MEM_FRAC_STATIC}
557- -e KV_CACHE_DTYPE=${KV_CACHE_DTYPE}
558- -e BLOCK_SIZE=${BLOCK_SIZE}
559- -e MAX_NUM_SEQS=${MAX_NUM_SEQS}
560- -e MAX_MODEL_LEN=${MAX_MODEL_LEN:- }
561- -e MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:- }
562- -e EXTRA_SERVER_ARGS=\$ {EXTRA_SERVER_ARGS:-}
563- -e IBDEVICES=${IBDEVICES:- }
564- )
565- elif [[ " $ENGINE " == " tilert" ]]; then
544+ if [[ " $ENGINE " == " tilert" ]]; then
566545 DOCKER_ENV_ENGINE=(
567546 -e MODEL_PATH=$DOCKER_MODEL_PATH
568547 -e PREFILL_IMAGE=${PREFILL_IMAGE}
@@ -671,84 +650,6 @@ echo \"Rank \$SLURM_PROCID on \$(hostname)\"
671650eval \"\$ DOCKER_CMD_DETECT\"
672651echo \" [docker-detect] rank \$ SLURM_PROCID: DOCKER_CMD=\$ DOCKER_CMD\"
673652
674- # Enable out-of-tree RDMA library mounts for atom-disagg (mooncake requires host RDMA stack)
675- RDMA_MOUNTS=()
676- if [[ " $ENGINE " == " atom-disagg" ]]; then
677-
678- # When the container base OS differs from the host (e.g. Ubuntu 24.04 image
679- # on a 22.04 host), the container's bundled libibverbs/libionic may be
680- # ABI-incompatible with the host kernel drivers. Detect the NIC type and
681- # bind-mount the host's out-of-tree RDMA userspace libraries into the
682- # container so the RDMA stack always matches the running kernel.
683- _detect_nic_type() {
684- if [[ -n \"\$ {MORI_NIC_TYPE:-}\" ]]; then echo \"\$ MORI_NIC_TYPE\" ; return; fi
685- local bnxt=0 mlx5=0 ionic=0
686- if [[ -d /sys/class/infiniband ]]; then
687- for dev in /sys/class/infiniband/*; do
688- local name; name=\$ (basename \"\$ dev\" )
689- case \"\$ name\" in
690- bnxt_re*) ((bnxt++)) ;; mlx5*) ((mlx5++)) ;; ionic*) ((ionic++)) ;;
691- *)
692- local drv; drv=\$ (basename \"\$ (readlink -f \"\$ dev/device/driver\" 2>/dev/null)\" 2>/dev/null || true)
693- case \"\$ drv\" in bnxt*) ((bnxt++)) ;; mlx5*) ((mlx5++)) ;; ionic*) ((ionic++)) ;; esac ;;
694- esac
695- done
696- fi
697- if (( bnxt >= mlx5 && bnxt >= ionic && bnxt > 0 )); then echo bnxt
698- elif (( ionic >= mlx5 && ionic > 0 )); then echo ionic
699- else echo mlx5; fi
700- }
701-
702- _find_host_ibverbs() {
703- for c in /usr/lib64/libibverbs.so.1 /lib/x86_64-linux-gnu/libibverbs.so.1 /usr/lib/x86_64-linux-gnu/libibverbs.so.1.14.39.0 /usr/lib/x86_64-linux-gnu/libibverbs.so.1; do
704- local r; r=\$ (readlink -f \"\$ c\" 2>/dev/null || true)
705- [[ \"\$ r\" == *libibverbs.so.1.14.57.0 ]] && continue
706- if [[ -f \"\$ r\" ]]; then echo \"\$ r\" ; return; fi
707- done
708- }
709-
710- _NIC_TYPE=\$ (_detect_nic_type)
711- echo \" [rdma] NIC type: \$ {_NIC_TYPE} on \$ (hostname)\"
712-
713- if [[ \"\$ _NIC_TYPE\" == \" ionic\" || \"\$ _NIC_TYPE\" == \" bnxt\" ]]; then
714- _host_ibv=\$ (_find_host_ibverbs)
715- if [[ -n \"\$ _host_ibv\" ]]; then
716- RDMA_MOUNTS+=(-v \"\$ _host_ibv:/lib/x86_64-linux-gnu/libibverbs.so.1\" )
717- fi
718- fi
719-
720- if [[ \"\$ _NIC_TYPE\" == \" ionic\" ]]; then
721- for _dir in /usr/local/lib /usr/lib/x86_64-linux-gnu; do
722- for _lib in \"\$ _dir\" /libionic*.so; do
723- [[ -f \"\$ _lib\" ]] || continue
724- _real=\$ (readlink -f \"\$ _lib\" )
725- [[ -f \"\$ _real\" ]] && RDMA_MOUNTS+=(-v \"\$ _real:\$ _real\" )
726- RDMA_MOUNTS+=(-v \"\$ _lib:/usr/lib/x86_64-linux-gnu/\$ (basename \"\$ _lib\" )\" )
727- done
728- done
729- if [[ -d /usr/lib/x86_64-linux-gnu/libibverbs ]]; then
730- for _lib in /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav*.so; do
731- [[ -f \"\$ _lib\" ]] && RDMA_MOUNTS+=(-v \"\$ _lib:\$ _lib\" )
732- done
733- fi
734- [[ -d /etc/libibverbs.d ]] && RDMA_MOUNTS+=(-v /etc/libibverbs.d:/etc/libibverbs.d:ro)
735- elif [[ \"\$ _NIC_TYPE\" == \" bnxt\" ]]; then
736- for _lib in /usr/local/lib/libbnxt_re-rdmav*.so; do
737- [[ -f \"\$ _lib\" ]] && RDMA_MOUNTS+=(-v \"\$ _lib:/usr/lib/x86_64-linux-gnu/libibverbs/\$ (basename \"\$ _lib\" )\" )
738- done
739- for _lib in /usr/local/lib/libbnxt_re.so; do
740- [[ -f \"\$ _lib\" ]] && RDMA_MOUNTS+=(-v \"\$ _lib:/usr/lib/x86_64-linux-gnu/\$ (basename \"\$ _lib\" )\" )
741- done
742- [[ -d /etc/libibverbs.d ]] && RDMA_MOUNTS+=(-v /etc/libibverbs.d:/etc/libibverbs.d:ro)
743- fi
744-
745- if [[ \$ {#RDMA_MOUNTS[@]} -gt 0 ]]; then
746- echo \" [rdma] bind-mounts: \$ {RDMA_MOUNTS[*]}\"
747- else
748- echo \" [rdma] no out-of-tree RDMA mounts needed\"
749- fi
750- fi # end: if ENGINE == atom-disagg
751-
752653RANK_IMAGE=
753654if [[ \" $ENGINE \" == \" tilert\" && \"\$ SLURM_PROCID\" -lt \" $xP \" ]]; then
754655 RANK_IMAGE=\" $PREFILL_IMAGE \"
@@ -790,7 +691,6 @@ exec \$DOCKER_CMD run \
790691 -v ${HICACHE_MC_CONFIG} :/config/hicache_mc.env:ro \
791692 ${EXTRA_DOCKER_MOUNTS:- } \
792693 ${CLIENT_DOCKER_MOUNTS} \
793- \$ {RDMA_MOUNTS[@]+" \$ {RDMA_MOUNTS[@]}" } \
794694 ${DOCKER_ENV_COMMON[*]} \
795695 ${DOCKER_ENV_ENGINE[*]} \
796696 ${CLIENT_DOCKER_ENV} \
0 commit comments