#!/bin/bash
cd ~
DOCKER_IMAGE="egs-registry.cn-hangzhou.cr.aliyuncs.com/egs/vllm:0.17.1-pytorch2.10-cu128-20260312"
MODEL_PATH=${MODEL_PATH}

# Validate node type
if [ "${NODE_TYPE}" != "--head" ] && [ "${NODE_TYPE}" != "--worker" ]; then
    echo "Error: Node type must be --head or --worker"
    exit 1
fi

container_id=$(docker ps -aqf "name=node")
if [ -n "$container_id" ]; then
  docker rm -f "$container_id"
fi

# Command setup for head or worker node
RAY_START_CMD="ray start --block"
if [ "${NODE_TYPE}" == "--head" ]; then
    RAY_START_CMD+=" --head --port=6379"
else
    RAY_START_CMD+=" --address=${HEAD_NODE_ADDRESS}:6379"
fi

# Run the docker command with the user specified parameters and additional arguments
docker run -d \
    --entrypoint /bin/bash \
    --network host \
    --name node \
    --ipc host \
    -e NCCL_DEBUG=TRACE \
    -e NCCL_SOCKET_IFNAME=eth0,eth1 \
    -e NCCL_IB_DISABLE=0 \
    --privileged \
    --pids-limit 65535 \
    -v /etc/localtime:/etc/localtime:ro \
    -v "${MODEL_PATH}:/workspace/DeepSeek-V4" \
    -e VLLM_HOST_IP=${NODE_ADDRESS} \
    "${DOCKER_IMAGE}"  -c "${RAY_START_CMD}"