Skip to content

Commit 24e0d82

Browse files
windreamerwindreamer
committed
build(docker): add jetson docker image
1 parent c93e764 commit 24e0d82

File tree

4 files changed

+140
-22
lines changed

4 files changed

+140
-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
@@ -146,28 +146,71 @@ set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}")
146146
set(CMAKE_CXX_STANDARD 17)
147147
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wall -ldl") # -Xptxas -v
148148

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

172215
message(STATUS "Building with CUDA archs: ${CMAKE_CUDA_ARCHITECTURES}")
173216

‎docker/Dockerfile.jetson‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Base images
2+
FROM nvcr.io/nvidia/l4t-base:r36.2.0
3+
ENV CUDA_VERSION=cu12.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_VERSION | sed 's/cu//g' | 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 git libopenblas-dev python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
20+
apt-get clean -y && \
21+
rm -rf /var/lib/apt/lists/* && \
22+
python${PYTHON_VERSION} -m venv /opt/py3 && \
23+
mkdir -p /wheels
24+
25+
# Should be in the lmdeploy root directory when building docker image
26+
COPY . /opt/lmdeploy
27+
WORKDIR /opt/lmdeploy
28+
29+
RUN --mount=type=cache,target=/root/.cache \
30+
--mount=type=cache,target=/opt/pytorch \
31+
pip install build change-wheel-version && \
32+
python -m build -w -o /wheels -v . && \
33+
change_wheel_version --local-version cu126 --delete-old-wheel /wheels/lmdeploy*.whl && \
34+
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
@@ -21,5 +21,5 @@ tiktoken
2121
torch<=2.8.0,>=2.0.0
2222
torchvision<=0.23.0,>=0.15.0
2323
transformers
24-
triton<=3.3.1,>=3.0.0; sys_platform == "linux"
24+
triton<=3.3.1,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
2525
uvicorn

0 commit comments

Comments
 (0)