Skip to content

Commit 52ce7bd

Browse files
authored
Merge pull request #619 from RedisAI/OnnxRuntime_set_allocator
Onnx runtime set allocator
2 parents 110030e + f27405b commit 52ce7bd

File tree

14 files changed

+485
-286
lines changed

14 files changed

+485
-286
lines changed

get_deps.sh

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
1212
if [[ $1 == --help || $1 == help ]]; then
1313
cat <<-END
1414
[ARGVARS...] get_deps.sh [cpu|gpu] [--help|help]
15-
15+
1616
Argument variables:
1717
CPU=1 Get CPU dependencies
1818
GPU=1 Get GPU dependencies
@@ -103,7 +103,7 @@ if [[ $WITH_TF != 0 ]]; then
103103
fi
104104
if [[ $ARCH == x64 ]]; then
105105
TF_ARCH=x86_64
106-
106+
107107
LIBTF_URL_BASE=https://storage.googleapis.com/tensorflow/libtensorflow
108108
elif [[ $ARCH == arm64v8 ]]; then
109109
TF_ARCH=arm64
@@ -134,7 +134,7 @@ if [[ $WITH_TF != 0 ]]; then
134134
mkdir $LIBTENSORFLOW.x
135135
tar xf $LIBTF_ARCHIVE --no-same-owner -C $LIBTENSORFLOW.x
136136
mv $LIBTENSORFLOW.x $LIBTENSORFLOW
137-
137+
138138
echo "Done."
139139
else
140140
echo "TensorFlow is in place."
@@ -152,7 +152,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
152152

153153
if [[ ! -d $LIBTFLITE ]]; then
154154
echo "Installing TensorFlow Lite ..."
155-
155+
156156
LIBTF_URL_BASE=https://s3.amazonaws.com/redismodules/tensorflow
157157
if [[ $OS == linux ]]; then
158158
TFLITE_OS="linux"
@@ -183,7 +183,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
183183
mkdir $LIBTFLITE.x
184184
tar xf $LIBTFLITE_ARCHIVE --no-same-owner -C $LIBTFLITE.x
185185
mv $LIBTFLITE.x $LIBTFLITE
186-
186+
187187
echo "Done."
188188
else
189189
echo "TensorFlow Lite is in place."
@@ -241,14 +241,14 @@ if [[ $WITH_PT != 0 ]]; then
241241

242242
[[ ! -f $LIBTORCH_ARCHIVE || $FORCE == 1 ]] && wget -q $LIBTORCH_URL
243243
fi
244-
244+
245245
rm -rf $LIBTORCH.x
246246
mkdir $LIBTORCH.x
247247

248248
tar xf $LIBTORCH_ARCHIVE --no-same-owner -C $LIBTORCH.x
249249
mv $LIBTORCH.x/libtorch $LIBTORCH
250250
rmdir $LIBTORCH.x
251-
251+
252252
echo "Done."
253253
else
254254
echo "libtorch is in place."
@@ -282,7 +282,7 @@ fi
282282

283283
################################################################################### ONNXRUNTIME
284284

285-
ORT_VERSION="1.6.0"
285+
ORT_VERSION="1.7.1"
286286

287287
if [[ $WITH_ORT != 0 ]]; then
288288
[[ $FORCE == 1 ]] && rm -rf $ONNXRUNTIME
@@ -297,15 +297,13 @@ if [[ $WITH_ORT != 0 ]]; then
297297
else
298298
ORT_BUILD="-gpu"
299299
fi
300+
ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
300301
if [[ $ARCH == x64 ]]; then
301302
ORT_ARCH=x64
302-
ORT_URL_BASE=https://github.com/microsoft/onnxruntime/releases/download/v${ORT_VERSION}
303303
elif [[ $ARCH == arm64v8 ]]; then
304304
ORT_ARCH=arm64
305-
ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
306305
elif [[ $ARCH == arm32v7 ]]; then
307306
ORT_ARCH=arm
308-
ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
309307
fi
310308
elif [[ $OS == macos ]]; then
311309
ORT_OS=osx
@@ -322,7 +320,7 @@ if [[ $WITH_ORT != 0 ]]; then
322320
mkdir $ONNXRUNTIME.x
323321
tar xzf ${ORT_ARCHIVE} --no-same-owner --strip-components=1 -C $ONNXRUNTIME.x
324322
mv $ONNXRUNTIME.x $ONNXRUNTIME
325-
323+
326324
echo "Done."
327325
else
328326
echo "ONNXRuntime is in place."

opt/build/onnxruntime/Dockerfile.arm7

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
44
ARG ONNXRUNTIME_BRANCH=rel-1.0.0
55
ARG ONNXRUNTIME_VER=1.0.0
66
ARG ARCH_FLAG="--arm"
7+
ARG ARCH=arm64v7
78

89
RUN [ "cross-build-start" ]
910

opt/build/onnxruntime/Dockerfile.x64

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,30 +5,36 @@ ARG OS=debian:buster
55
#----------------------------------------------------------------------------------------------
66
FROM ${OS}
77

8-
ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
9-
ARG ONNXRUNTIME_VER=1.0.0
8+
ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
9+
ARG ONNXRUNTIME_VER=1.7.1
10+
ARG ARCH=x64
1011

1112
RUN apt-get -qq update
12-
RUN apt-get -qq install -y curl wget tar git
13-
RUN apt-get -qq install -y build-essential cmake
14-
RUN apt-get -qq install -y libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev
15-
16-
RUN apt-get -qq install -y python3 python3-pip python3-dev
13+
RUN apt-get -qq install -y curl wget tar git \
14+
build-essential cmake \
15+
libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev \
16+
python3 python3-pip python3-dev python3-numpy
1717
RUN pip3 install --upgrade pip setuptools wheel
18-
# RUN pip3 install numpy
19-
RUN apt-get -q install -y python3-numpy
18+
19+
ENV LANG=en_US.UTF-8
20+
RUN apt-get install -y locales && \
21+
sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
22+
dpkg-reconfigure --frontend=noninteractive locales && \
23+
update-locale LANG=$LANG
2024

2125
WORKDIR /build
2226

23-
ADD ./pack.sh /build/
2427
ARG BUILDTYPE=MinSizeRel
28+
2529
ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
2630

27-
RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
28-
29-
RUN cd onnxruntime ;\
30-
./build.sh ${BUILDARGS} --update --build ;\
31-
./build.sh ${BUILDARGS} --build_shared_lib
31+
RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
32+
WORKDIR /build/onnxruntime
33+
RUN git fetch --recurse-submodules -j4
34+
RUN ./build.sh ${BUILDARGS} --update --build
35+
RUN ./build.sh ${BUILDARGS} --build_shared_lib
3236
# RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
3337

34-
RUN ./pack.sh ${ONNXRUNTIME_VER}
38+
ADD ./pack.sh /build
39+
WORKDIR /build
40+
RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
ARG OS=ubuntu18.04
2+
ARG CUDA_VER=11.0-cudnn8
3+
4+
FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
5+
6+
ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
7+
ARG ONNXRUNTIME_VER=1.7.1
8+
ARG ARCH=x64-gpu
9+
10+
RUN apt-get -qq update
11+
RUN apt-get -qq install -y curl \
12+
wget tar git build-essential \
13+
libcurl4-openssl-dev libssl-dev \
14+
libatlas-base-dev zlib1g-dev python3 python3-pip \
15+
python3-dev python3-numpy rsync
16+
17+
RUN pip3 install --upgrade pip setuptools wheel
18+
19+
RUN wget -q https://github.com/Kitware/CMake/releases/download/v3.19.5/cmake-3.19.5-Linux-x86_64.tar.gz -O /tmp/cmake.tgz
20+
WORKDIR /tmp
21+
RUN tar -xpf cmake.tgz
22+
RUN rsync -aqH cmake*/* /usr
23+
24+
25+
ENV LANG=en_US.UTF-8
26+
RUN apt-get install -y locales && \
27+
sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
28+
dpkg-reconfigure --frontend=noninteractive locales && \
29+
update-locale LANG=$LANG
30+
31+
WORKDIR /build
32+
33+
ARG BUILDTYPE=MinSizeRel
34+
ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
35+
36+
RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
37+
WORKDIR /build/onnxruntime
38+
RUN git fetch --recurse-submodules -j4
39+
RUN ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/cuda --cuda_home /usr/local/cuda --build_shared_lib --parallel
40+
41+
ADD ./pack.sh /build/
42+
WORKDIR /build
43+
RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}

opt/build/onnxruntime/Makefile

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11

22
ROOT=.
33

4-
VERSION ?= 1.0.0
4+
VERSION ?= 1.7.1
55
OSNICK ?= buster
66

77
#----------------------------------------------------------------------------------------------
@@ -16,7 +16,7 @@ DOCKER_OS.bionic=ubuntu:bionic
1616
DOCKER_OS.stretch=debian:stretch-slim
1717
DOCKER_OS.buster=debian:buster-slim
1818
DOCKER_OS=$(DOCKER_OS.$(OSNICK))
19-
19+
2020
#----------------------------------------------------------------------------------------------
2121

2222
define targets # (1=OP, 2=op)
@@ -40,10 +40,23 @@ IID_$(1)=$(1)_$(VERSION).iid
4040
CID_$(1)=$(1)_$(VERSION).cid
4141

4242
build_x64:
43-
@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64 \
43+
@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-x64:$(VERSION) -f Dockerfile.x64 \
4444
--build-arg OS=$(DOCKER_OS) $(ROOT)
4545
@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
46-
@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(2)-$(VERSION).tgz .
46+
@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-x64-$(VERSION).tgz .
47+
48+
.PHONY: build_x64
49+
endef
50+
51+
define build_x64-gpu # (1=arch, 2=tar-arch)
52+
IID_$(1)=$(1)_$(VERSION).iid
53+
CID_$(1)=$(1)_$(VERSION).cid
54+
55+
build_x64-gpu:
56+
@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64-gpu \
57+
--build-arg OS=$(DOCKER_OS) $(ROOT)
58+
@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
59+
@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(1)-$(VERSION).tgz .
4760

4861
.PHONY: build_x64
4962
endef
@@ -68,11 +81,18 @@ endef
6881

6982
define publish_x64 # (1=arch, 2=tar-arch)
7083
publish_x64:
71-
@aws s3 cp $(STEM)-$(2)-$(VERSION).tar.gz s3://$(S3_URL)/ --acl public-read
84+
@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
7285

7386
.PHONY: publish_x64
7487
endef
7588

89+
define publish_x64-gpu # (1=arch, 2=tar-arch)
90+
publish_x64-gpu:
91+
@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
92+
93+
.PHONY: publish_x64-gpu
94+
endef
95+
7696
define publish_arm # (1=arch, 2=tar-arch)
7797
publish_$(1):
7898
@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
@@ -90,6 +110,7 @@ all: build publish
90110
build: $(BUILD_TARGETS)
91111

92112
$(eval $(call build_x64,x64,x86_64))
113+
$(eval $(call build_x64-gpu,x64-gpu,x86_64-gpu))
93114
$(eval $(call build_arm,arm64v8,arm64))
94115
$(eval $(call build_arm,arm32v7,arm))
95116

@@ -101,7 +122,8 @@ endif
101122

102123
publish: $(PUBLISH_TARGETS)
103124

104-
$(eval $(call publish_x64,x64,x86_64))
125+
$(eval $(call publish_x64,x64,x64))
126+
$(eval $(call publish_x64-gpu,x64-gpu,x64-gpu))
105127
$(eval $(call publish_arm,arm64v8,arm64))
106128
$(eval $(call publish_arm,arm32v7,arm))
107129

opt/build/onnxruntime/pack.sh

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,11 @@
22

33
set -e
44
VER="$1"
5+
PLATFORM="$2"
56

67
mkdir -p pack/include pack/lib
78
cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
8-
cp onnxruntime/docs/C_API.md pack/
9+
cp onnxruntime/docs/C_API_Guidelines.md pack/
910
cp onnxruntime/LICENSE pack/
1011
cp onnxruntime/README.md pack/
1112
cp onnxruntime/ThirdPartyNotices.txt pack/
@@ -14,8 +15,9 @@ cd onnxruntime/
1415
git rev-parse HEAD > ../pack/GIT_COMMIT_ID
1516
cd ..
1617
cp onnxruntime/include/onnxruntime/core/session/onnxruntime_c_api.h pack/include/
18+
cp onnxruntime/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h pack/include/
1719
cd pack/lib/
1820
ln -s libonnxruntime.so.${VER} libonnxruntime.so
1921
cd ../..
20-
mv pack onnxruntime-linux-arm64-${VER}
21-
tar czf onnxruntime-linux-arm64-${VER}.tgz onnxruntime-linux-arm64-${VER}/
22+
mv pack onnxruntime-linux-${PLATFORM}-${VER}
23+
tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/

opt/redis_valgrind.sup

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,6 @@
1212
obj:*/libtensorflow_framework.so.*
1313
}
1414

15-
{
16-
ignore_unversioned_libs
17-
Memcheck:Leak
18-
...
19-
obj:*/libonnxruntime.so.*
20-
}
21-
2215
{
2316
ignore_unversioned_libs
2417
Memcheck:Leak

src/backends.c

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -451,10 +451,27 @@ int RAI_LoadBackend_ONNXRuntime(RedisModuleCtx *ctx, const char *path) {
451451
return REDISMODULE_ERR;
452452
}
453453

454-
RAI_backends.onnx = backend;
454+
backend.get_memory_info =
455+
(unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryInfoORT");
456+
if (backend.get_memory_info == NULL) {
457+
dlclose(handle);
458+
RedisModule_Log(ctx, "warning",
459+
"Backend does not export RAI_GetMemoryInfoORT. ONNX backend "
460+
"not loaded from %s",
461+
path);
462+
}
463+
backend.get_memory_access_num =
464+
(unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryAccessORT");
465+
if (backend.get_memory_access_num == NULL) {
466+
dlclose(handle);
467+
RedisModule_Log(ctx, "warning",
468+
"Backend does not export RAI_GetMemoryAccessORT. ONNX backend "
469+
"not loaded from %s",
470+
path);
471+
}
455472

473+
RAI_backends.onnx = backend;
456474
RedisModule_Log(ctx, "notice", "ONNX backend loaded from %s", path);
457-
458475
return REDISMODULE_OK;
459476
}
460477

src/backends.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,11 @@ typedef struct RAI_LoadedBackend {
7777
// Returns the backend version.
7878
const char *(*get_version)(void);
7979

80+
// Returns the backend's memory usage for INFO report
81+
unsigned long long (*get_memory_info)(void);
82+
83+
// Returns the number of times that Redis accessed backend allocator.
84+
unsigned long long (*get_memory_access_num)(void);
8085
} RAI_LoadedBackend;
8186

8287
typedef struct RAI_LoadedBackends {

0 commit comments

Comments
 (0)