Skip to content

Commit 2565ed1

Browse files
windreamerwindreamer
committed
build(docker): add jetson docker image
1 parent 9b01105 commit 2565ed1

File tree

4 files changed

+141
-22
lines changed

4 files changed

+141
-22
lines changed

‎.github/workflows/test_docker.yml‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,44 @@ jobs:
103103
with:
104104
image: lmdeploy:ascend
105105
github-token: ${{ secrets.GITHUB_TOKEN }}
106+
107+
test_jetson_docker_image:
108+
permissions:
109+
pull-requests: write
110+
runs-on: ubuntu-22.04-arm
111+
steps:
112+
- name: Checkout repository
113+
uses: actions/checkout@v3
114+
with:
115+
ref: ${{github.event.inputs.repo_ref}}
116+
- name: Free disk space
117+
uses: jlumbroso/free-disk-space@main
118+
with:
119+
# This might remove tools that are actually needed, if set to "true" but frees about 6 GB
120+
tool-cache: false
121+
docker-images: false
122+
# All of these default to true, but feel free to set to "false" if necessary for your workflow
123+
android: true
124+
dotnet: true
125+
haskell: true
126+
large-packages: true
127+
swap-storage: false
128+
- name: Set up Docker Buildx
129+
uses: docker/setup-buildx-action@v3
130+
- name: Get docker info
131+
run: |
132+
docker info
133+
# remove http extraheader
134+
git config --local --unset "http.https://github.com/.extraheader"
135+
- name: Build Docker image
136+
run: |
137+
docker build . -t lmdeploy:jetson -f docker/Dockerfile.jetson
138+
- name: Test image with lmdeploy check_env
139+
run: |
140+
docker images
141+
docker run --rm lmdeploy:jetson lmdeploy check_env
142+
- name: Dive
143+
uses: MaxymVlasov/[email protected]
144+
with:
145+
image: lmdeploy:jetson
146+
github-token: ${{ secrets.GITHUB_TOKEN }}

‎CMakeLists.txt‎

Lines changed: 64 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -148,28 +148,71 @@ set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}")
148148
set(CMAKE_CXX_STANDARD 17)
149149
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wall -ldl") # -Xptxas -v
150150

151-
# TODO: build for sm_72 & sm_87 on aarch64 platform (Jetson devices)
152-
if (NOT CMAKE_CUDA_ARCHITECTURES)
153-
set(CMAKE_CUDA_ARCHITECTURES 70-real 75-real) # V100, 2080
154-
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11")
155-
list(APPEND CMAKE_CUDA_ARCHITECTURES 80-real) # A100
156-
endif ()
157-
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.1")
158-
list(APPEND CMAKE_CUDA_ARCHITECTURES 86-real) # 3090
159-
endif ()
160-
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.8")
161-
list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real) # 4090
162-
endif ()
163-
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.0")
164-
list(APPEND CMAKE_CUDA_ARCHITECTURES 90a-real) # H100
165-
endif ()
166-
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.8")
167-
list(APPEND CMAKE_CUDA_ARCHITECTURES 120a-real) # 5090
168-
endif ()
169-
if (MSVC)
170-
list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES 80-real 90a-real)
151+
if(CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64")
152+
set(ARCH "x86_64")
153+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "amd64")
154+
set(ARCH "x86_64")
155+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "AMD64")
156+
# cmake reports AMD64 on Windows, but we might be building for 32-bit.
157+
if(CMAKE_SIZEOF_VOID_P EQUAL 8)
158+
set(ARCH "x86_64")
159+
else()
160+
set(ARCH "x86")
161+
endif()
162+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "x86")
163+
set(ARCH "x86")
164+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "i386")
165+
set(ARCH "x86")
166+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "i686")
167+
set(ARCH "x86")
168+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
169+
set(ARCH "aarch64")
170+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
171+
set(ARCH "aarch64")
172+
# Apple A12 Bionic chipset which is added in iPhone XS/XS Max/XR uses arm64e architecture.
173+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64e")
174+
set(ARCH "aarch64")
175+
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^arm*")
176+
set(ARCH "arm")
177+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "mips")
178+
# Just to avoid the “unknown processor” error.
179+
set(ARCH "generic")
180+
elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "ppc64le")
181+
set(ARCH "ppc64le")
182+
else()
183+
message(FATAL_ERROR "Unknown processor:" ${CMAKE_SYSTEM_PROCESSOR})
184+
endif()
185+
186+
187+
if(ARCH STREQUAL "x86_64")
188+
if (NOT CMAKE_CUDA_ARCHITECTURES)
189+
set(CMAKE_CUDA_ARCHITECTURES 70-real 75-real) # V100, 2080
190+
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11")
191+
list(APPEND CMAKE_CUDA_ARCHITECTURES 80-real) # A100
192+
endif ()
193+
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.1")
194+
list(APPEND CMAKE_CUDA_ARCHITECTURES 86-real) # 3090
195+
endif ()
196+
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.8")
197+
list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real) # 4090
198+
endif ()
199+
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.0")
200+
list(APPEND CMAKE_CUDA_ARCHITECTURES 90a-real) # H100
201+
endif ()
202+
if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.8")
203+
list(APPEND CMAKE_CUDA_ARCHITECTURES 120a-real) # 5090
204+
endif ()
205+
if (MSVC)
206+
list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES 80-real 90a-real)
207+
endif ()
171208
endif ()
172-
endif ()
209+
elseif(ARCH STREQUAL "aarch64")
210+
if (NOT CMAKE_CUDA_ARCHITECTURES)
211+
set(CMAKE_CUDA_ARCHITECTURES 72-real 87-real) # Jetson
212+
endif()
213+
else()
214+
message(FATAL_ERROR "Unsupported Architecture:" ${ARCH})
215+
endif()
173216

174217
message(STATUS "Building with CUDA archs: ${CMAKE_CUDA_ARCHITECTURES}")
175218

‎docker/Dockerfile.jetson‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Base images
2+
FROM nvcr.io/nvidia/l4t-base:r36.2.0
3+
ENV CUDA_VER=12.6 \
4+
PYTHON_VERSION=3.10 \
5+
PATH=/opt/py3/bin:/root/.local/bin:/usr/local/cuda/bin:${PATH}
6+
7+
RUN --mount=type=cache,target=/root/.cache \
8+
--mount=type=cache,target=/tmp/download \
9+
export CUDA_SUFFIX=$(echo $CUDA_VER | sed 's/\./-/g') && \
10+
cd /tmp/download && \
11+
mkdir -p /opt/nvidia/l4t-packages/ && \
12+
touch /opt/nvidia/l4t-packages/.nv-l4t-disable-boot-fw-update-in-preinstall && \
13+
wget -q "https://repo.download.nvidia.com/jetson/t234/pool/main/n/nvidia-l4t-core/nvidia-l4t-core_36.2.0-20231218214829_arm64.deb" && \
14+
wget -q "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/arm64/cuda-keyring_1.1-1_all.deb" && \
15+
yes | dpkg -i nvidia-l4t-core_*.deb cuda-keyring_*.deb && \
16+
rm -rf *.deb *.deb.* && \
17+
apt update -y && \
18+
apt-get install -y --no-install-recommends \
19+
cuda-toolkit-${CUDA_SUFFIX} cuda-compat-${CUDA_SUFFIX} libcudnn9-cuda-12 libcusparselt0 cudss \
20+
git libopenblas-dev python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
21+
apt-get clean -y && \
22+
rm -rf /var/lib/apt/lists/* && \
23+
python${PYTHON_VERSION} -m venv /opt/py3 && \
24+
mkdir -p /wheels
25+
26+
# Should be in the lmdeploy root directory when building docker image
27+
COPY . /opt/lmdeploy
28+
WORKDIR /opt/lmdeploy
29+
30+
RUN --mount=type=cache,target=/root/.cache \
31+
--mount=type=cache,target=/opt/pytorch \
32+
pip install build change-wheel-version && \
33+
python -m build -w -o /wheels -v . && \
34+
change_wheel_version --local-version cu126 --delete-old-wheel /wheels/lmdeploy*.whl && \
35+
pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/

‎requirements/runtime_cuda.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,5 @@ tiktoken
2323
torch<=2.8.0,>=2.0.0
2424
torchvision<=0.23.0,>=0.15.0
2525
transformers
26-
triton<=3.4.0,>=3.0.0; sys_platform == "linux"
26+
triton<=3.4.0,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
2727
uvicorn

0 commit comments

Comments
 (0)