NVIDIA Brev GPU Launchable to test CUDA Backend

[Maalvi14]

Hello everyone,

I've been working on some features for the CUDA Backend, i took the time to make Brev launchable (great part of the script is based on the circle-ci config which deploy the gu instances for the PRs) for everyone to deploy and run tests on.

The script is still not fully working, but I wanted to share in nonetheless so anyone who wants to hop in and improve can do so. :)

Brev Launchable

---

Startup Script

#!/bin/bash
set -euo pipefail

# ================================
# MLX CUDA editable setup for Brev
# Tested on Ubuntu 22.04 / 24.04
# ================================

# ---- Config (override via env) ----
: "${PYTHON_VER:=python3.10}"     # use system python; Brev images usually have 3.10/3.11
: "${VENV_DIR:=.venv}"            # venv location inside the repo
: "${MLX_REPO_DIR:=}"             # optional; set to existing clone path; if empty, will git clone
: "${MLX_GIT_URL:=https://github.com/ml-explore/mlx.git}"
: "${MLX_GIT_BRANCH:=main}"       # or your feature/PR branch

echo ">>> Checking GPU ..."
if ! command -v nvidia-smi >/dev/null 2>&1; then
  echo "ERROR: NVIDIA driver not visible. In Brev, pick a GPU image/instance. Aborting."
  exit 1
fi
nvidia-smi || true

# ---- Detect Ubuntu version ----
source /etc/os-release
UBU_VER="${VERSION_ID:-unknown}"
echo ">>> Detected Ubuntu ${UBU_VER}"

# ---- Base deps ----
echo ">>> Installing base packages ..."
sudo apt-get update -y
sudo apt-get install -y --no-install-recommends \
  build-essential git curl ca-certificates pkg-config \
  libblas-dev liblapack-dev liblapacke-dev \
  python3-venv ${PYTHON_VER} ${PYTHON_VER%-*}-dev

# ---- CUDA + cuDNN (match CI as close as practical) ----
# Prefer CUDA 12.9 on 24.04 (as in CI). Fall back to toolkit-12-9 meta-package if available.
echo ">>> Installing CUDA Toolkit and cuDNN ..."
if [[ "${UBU_VER}" == "24.04" ]]; then
  wget -q 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 apt-get update -y
  sudo apt-get install -y cuda-toolkit-12-9 libcudnn9-dev-cuda-12
elif [[ "${UBU_VER}" == "22.04" ]]; then
  wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
  sudo dpkg -i cuda-keyring_1.1-1_all.deb
  sudo apt-get update -y
  # 12.9 might be available; if not, install generic 12.x toolkit
  if apt-cache policy cuda-toolkit-12-9 | grep -q Candidate; then
    sudo apt-get install -y cuda-toolkit-12-9 libcudnn9-dev-cuda-12
  else
    sudo apt-get install -y cuda-toolkit-12 libcudnn9-dev-cuda-12
  fi
else
  echo "WARN: Unknown Ubuntu version ${UBU_VER}. Trying CUDA 12 generic ..."
  sudo apt-get install -y cuda-toolkit-12 libcudnn9-dev-cuda-12 || true
fi

# (Optional) NCCL — useful for multi-GPU tests; harmless if single-GPU
if apt-cache policy libnccl-dev | grep -q Candidate; then
  echo ">>> Installing NCCL ..."
  sudo apt-get install -y libnccl2 libnccl-dev
fi

# ---- CUDA env ----
echo ">>> Exporting CUDA env ..."
if ! grep -q "CUDA PATH export" ~/.bashrc 2>/dev/null; then
  {
    echo ""
    echo "# CUDA PATH export"
    echo 'export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}'
    echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}'
  } >> ~/.bashrc
fi
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
which nvcc
nvcc --version

# ---- uv (fast Python installer/runner) ----
if ! command -v uv >/dev/null 2>&1; then
  echo ">>> Installing uv ..."
  curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.local/bin:$PATH"
fi

# ---- Clone or use existing MLX repo ----
if [[ -z "${MLX_REPO_DIR}" ]]; then
  echo ">>> Cloning MLX repo ..."
  MLX_REPO_DIR="$HOME/mlx"
  if [[ -d "${MLX_REPO_DIR}" ]]; then
    echo "Repo dir ${MLX_REPO_DIR} already exists; using it."
  else
    git clone --recursive "${MLX_GIT_URL}" "${MLX_REPO_DIR}"
  fi
fi
cd "${MLX_REPO_DIR}"
git fetch --all
git checkout "${MLX_GIT_BRANCH}"
git submodule update --init --recursive

# ---- Python env + editable install (CUDA) ----
echo ">>> Creating venv and installing MLX in editable mode (CUDA) ..."
uv venv --python "${PYTHON_VER#python}"
# Ensure cmake present (nanobind build needs it)
uv pip install --upgrade pip
uv pip install cmake wheel setuptools

# Match CircleCI CUDA build flags:
#   -DMLX_BUILD_CUDA=ON
#   -DCMAKE_CUDA_COMPILER=`which nvcc`
DEBUG=1 \
CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=$(which nvcc)" \
  uv pip install -e ".[dev]" -v

# ---- Quick GPU smoke test ----
echo ">>> Running MLX GPU smoke test ..."
cat > /tmp/mlx_cuda_smoke.py <<'PY'
import mlx.core as mx

# Show available devices
print("Devices:", mx.devices())

# Try forcing GPU (if available) and run a tiny op
if hasattr(mx, "gpu"):
    try:
        mx.set_default_device(mx.gpu)
        a = mx.arange(0, 8, dtype=mx.float32)
        b = mx.sin(a)
        # materialize / sync
        mx.eval(b)
        print("GPU OK:", b)
    except Exception as e:
        print("GPU test failed:", e)
else:
    print("mx.gpu not found in this build; ensure MLX CUDA was enabled.")
PY

source "${VENV_DIR}/bin/activate"
python /tmp/mlx_cuda_smoke.py

echo ">>> Done. MLX is installed in EDITABLE mode with CUDA."
echo "Activate your env with:  source ${MLX_REPO_DIR}/${VENV_DIR}/bin/activate"