Upgrade VM image (#325)

Author: ultmaster

pyproject.toml

@@ -95,7 +95,8 @@ torch-stable = [
   "transformers>=4.55.0",
   # vLLM 0.11.1 requires PyTorch 2.9.0, which is incompatible with flash-attn
   # https://github.com/Dao-AILab/flash-attention/issues/1967
-  "vllm>=0.10.2,!=0.11.1",
+  # Similar issues with vLLM 0.11.2
+  "vllm>=0.10.2,!=0.11.1,!=0.11.2",
   # LiteLLM can then be upgraded with new vLLM
   "litellm[proxy]>=1.78",
 ]

scripts/build_vm_image.sh

@@ -25,7 +25,12 @@ sudo apt-get install -y \
     tmux \
     vim \
     git-lfs \
-    nodejs
+    nodejs \
+    gnupg2 \
+    apt-transport-https \
+    ca-certificates \
+    gnupg \
+    lsb-release
 
 git lfs install
 
@@ -37,9 +42,9 @@ sudo reboot now
 
 # Install CUDA Toolkit
 wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
-sudo dpkg -i cuda-keyring_1.1-1_all.deb
+sudo dpkg -i cuda-keyring_1.1-1_all.deb && rm cuda-keyring_1.1-1_all.deb
 sudo apt-get update
-sudo apt-get -y install cuda-toolkit-12-8
+sudo apt-get -y install cuda-toolkit
 sudo reboot now
 
 # Add paths globally
@@ -49,6 +54,67 @@ export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 EOF
 sudo chmod +x /etc/profile.d/cuda.sh
 
+# Add Docker's official GPG key
+sudo install -m 0755 -d /etc/apt/keyrings
+sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
+sudo chmod a+r /etc/apt/keyrings/docker.asc
+
+# Add the repository to Apt sources:
+sudo tee /etc/apt/sources.list.d/docker.sources <<EOF
+Types: deb
+URIs: https://download.docker.com/linux/ubuntu
+Suites: $(. /etc/os-release && echo "${UBUNTU_CODENAME:-$VERSION_CODENAME}")
+Components: stable
+Signed-By: /etc/apt/keyrings/docker.asc
+EOF
+
+sudo apt -y update
+
+# Install the Docker packages
+sudo apt -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
+
+# Create docker group only if it doesn't exist
+# sudo groupadd docker
+
+# Add current user to docker group if not already a member
+sudo usermod -aG docker "$USER"
+# A hack to add cloudtest user to docker group as well
+sudo sed -i '/^docker:/ s/$/,cloudtest/' /etc/group
+# This shouldn't be run on CI
+# newgrp docker
+
+# Install NVIDIA Container Toolkit
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+
+sudo apt-get update
+export NVIDIA_CONTAINER_TOOLKIT_VERSION=1.18.0-1
+sudo apt-get install -y \
+    nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}
+
+# Configure the NVIDIA Container Toolkit
+sudo nvidia-ctk runtime configure --runtime=docker
+sudo systemctl restart docker
+
+# Install Azure CLI
+curl -sLS https://packages.microsoft.com/keys/microsoft.asc |
+  gpg --dearmor | sudo tee /etc/apt/keyrings/microsoft.gpg > /dev/null
+sudo chmod go+r /etc/apt/keyrings/microsoft.gpg
+AZ_DIST=$(lsb_release -cs)
+echo "Types: deb
+URIs: https://packages.microsoft.com/repos/azure-cli/
+Suites: ${AZ_DIST}
+Components: main
+Architectures: $(dpkg --print-architecture)
+Signed-by: /etc/apt/keyrings/microsoft.gpg" | sudo tee /etc/apt/sources.list.d/azure-cli.sources
+sudo apt-get update
+sudo apt-get install -y azure-cli
+
 # Disable the periodical apt-get upgrade.
 # Sometimes, unattended upgrade blocks apt-get install
 sudo sed -i -e "s/Update-Package-Lists \"1\"/Update-Package-Lists \"0\"/g" /etc/apt/apt.conf.d/10periodic

scripts/litellm_sanity_check.py

@@ -8,7 +8,7 @@
 
 
 def main() -> None:
-    client = openai.OpenAI()
+    client = openai.OpenAI(timeout=30.0)
     models = client.models.list()
     print("Available models:", models)