misc: update devcontainer (#986)

Use cu128 and nightly pytorch for blackwell development.

Author: yzh119

.devcontainer/Dockerfile

@@ -1,31 +1,24 @@
-FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
+FROM nvidia/cuda:12.8.0-devel-ubuntu24.04
 
 # Update package lists and install system dependencies
 RUN apt-get update && apt-get install -y \
     curl \
+    wget \
     git \
     clang-format \
-    libibverbs-dev \
-    librdmacm-dev \
-    rdma-core \
-    libnuma-dev \
+    clangd-19 \
     vim \
-    openmpi-bin \
-    libopenmpi-dev \
     zsh \
     && rm -rf /var/lib/apt/lists/*
 
 # Install oh-my-zsh
 RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
 
-# Install powerlevel10k theme
-RUN git clone --depth=1 https://github.com/romkatv/powerlevel10k.git ${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}/themes/powerlevel10k
-
 # Install zsh-autosuggestions
 RUN git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions
 
 # Configure zsh
-RUN sed -i 's/ZSH_THEME="robbyrussell"/ZSH_THEME="powerlevel10k\/powerlevel10k"/' ~/.zshrc && \
+RUN sed -i 's/ZSH_THEME="robbyrussell"/ZSH_THEME="fino-time"/' ~/.zshrc && \
     sed -i 's/plugins=(git)/plugins=(git zsh-autosuggestions)/' ~/.zshrc
 
 # Create a non-root user
@@ -53,6 +46,21 @@ RUN cp -r /root/.oh-my-zsh /home/$USERNAME/.oh-my-zsh && \
 USER $USERNAME
 WORKDIR /home/$USERNAME
 
+# Install python
+COPY install/install_python.sh /install/install_python.sh
+RUN bash /install/install_python.sh py312
+
+# clangd
+ENV PATH="/usr/lib/llvm-19/bin:$PATH"
+# conda
+ENV PATH="/home/devuser/conda/bin:$PATH"
+ENV PATH="/home/devuser/envs/py312/bin:$PATH"
+
+# Install python packages
+COPY install/install_python_packages.sh /install/install_python_packages.sh
+RUN bash /install/install_python_packages.sh
+RUN echo "source activate py312" >> /home/devuser/.zshrc
+
 # Set zsh as default shell
 ENV SHELL=/bin/zsh
 CMD [ "zsh" ]

.devcontainer/devcontainer.json

@@ -10,8 +10,15 @@
     "customizations": {
         "vscode": {
             "extensions": [
+                "llvm-vs-code-extensions.vscode-clangd",
+                "ms-python.python",
+                "ms-python.black-formatter",
+                "nvidia.nsight-vscode-edition"
             ]
         }
     },
+    "mounts": [
+        "type=bind,source=${localEnv:HOME}/.ssh,target=/home/devuser/.ssh,readonly"
+    ],
     "remoteUser": "devuser"
 }

.devcontainer/install/install_python.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+set -e
+set -u
+set -o pipefail
+
+
+# Install python and pip. Don't modify this to add Python package dependencies,
+wget -O Miniforge3.sh "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
+bash Miniforge3.sh -b -p /home/devuser/conda
+
+/home/devuser/conda/bin/conda create -n $1 python=3.12

.devcontainer/install/install_python_packages.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+set -e
+set -u
+
+pip3 install ninja pytest numpy scipy build cuda-python pytest
+pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+pip3 install pre-commit