Nvidia Jetson support (phase 1) (#535)

Author: rafie

.dockerignore

@@ -1,5 +1,6 @@
 /bin/
 /deps/
+/docs/
 /install*
 /build/
 .venv/

Dockerfile.jetson

@@ -0,0 +1,99 @@
+# BUILD redisfab/redisai:${VERSION}-jetson-${ARCH}-${OSNICK}
+
+ARG REDIS_VER=6.0.9
+
+# OSNICK=bionic|centos7|centos6
+ARG OSNICK=bionic
+
+# OS=ubuntu18.04|ubuntu16.04|centos7
+ARG OS=ubuntu18.04
+
+# ARCH=arm64v8|arm32v7
+ARG ARCH=arm64v8
+
+ARG CUDA_VER=10.2-cudnn7
+
+ARG L4T_VER=r32.4.4
+
+ARG PACK=0
+ARG TEST=0
+
+#----------------------------------------------------------------------------------------------
+FROM redisfab/redis:${REDIS_VER}-${ARCH}-${OSNICK} AS redis
+FROM redisfab/jetpack:4.4.1-arm64v8-l4t as builder
+
+ARG OSNICK
+ARG OS
+ARG ARCH
+ARG REDIS_VER
+ARG CUDA_VER
+ARG L4T_VER
+
+RUN echo "Building for ${OSNICK} (${OS}) for ${ARCH} [with Redis ${REDIS_VER}]"
+
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+WORKDIR /build
+COPY --from=redis /usr/local/ /usr/local/
+
+COPY ./opt/ opt/
+COPY ./tests/flow/test_requirements.txt tests/flow/
+
+RUN PIP=1 FORCE=1 ./opt/readies/bin/getpy3
+RUN ./opt/system-setup.py
+
+ARG DEPS_ARGS="GPU=1 JETSON=1 WITH_PT=1 WITH_TF=1 WITH_TFLITE=0 WITH_ORT=0"
+COPY ./get_deps.sh .
+# RUN if [ "$DEPS_ARGS" = "" ]; then ./get_deps.sh gpu; else env $DEPS_ARGS ./get_deps.sh gpu; fi
+RUN set -e ;\
+    env $DEPS_ARGS ./get_deps.sh
+
+ARG BUILD_ARGS="GPU=1 JETSON=1 WITH_TF=1 WITH_PT=1 WITH_TFLITE=0 WITH_ORT=0"
+
+ADD ./ /build
+RUN bash -c "set -e ;\
+    . ./opt/readies/bin/sourced ./profile.d ;\
+    make -C opt build  $BUILD_ARGS SHOW=1"
+
+ARG PACK
+ARG TEST
+
+RUN mkdir -p bin/artifacts
+RUN set -e ;\
+    if [ "$PACK" = "1" ]; then make -C opt pack GPU=1; fi
+
+RUN set -e ;\
+    if [ "$TEST" = "1" ]; then \
+        TEST= make -C opt test GPU=1 $BUILD_ARGS NO_LFS=1 ;\
+        if [[ -d test/logs ]]; then \
+            tar -C test/logs -czf bin/artifacts/test-logs-gpu.tgz . ;\
+        fi ;\
+    fi
+
+#----------------------------------------------------------------------------------------------
+# FROM nvidia/cuda:${CUDA_VER}-runtime-${OS}
+# FROM nvidia/cuda-arm64:11.1-runtime-ubuntu18.04
+# FROM nvcr.io/nvidia/l4t-base:${L4T_VER}
+FROM redisfab/jetpack:4.4.1-arm64v8-l4t
+
+ARG OS
+ARG L4T_VER
+
+RUN if [ ! -z $(command -v apt-get) ]; then apt-get -qq update; apt-get -q install -y libgomp1; fi
+RUN if [ ! -z $(command -v yum) ]; then yum install -y libgomp; fi 
+
+ENV REDIS_MODULES /usr/lib/redis/modules
+RUN mkdir -p $REDIS_MODULES/
+
+COPY --from=redis /usr/local/ /usr/local/
+COPY --from=builder /build/install-gpu/ $REDIS_MODULES/
+
+COPY --from=builder /build/bin/artifacts/ /var/opt/redislabs/artifacts
+
+# COPY --from=builder /usr/local/etc/redis/redis.conf /usr/local/etc/redis/
+
+WORKDIR /data
+EXPOSE 6379
+# CMD ["/usr/local/bin/redis-server", "/usr/local/etc/redis/redis.conf", "--loadmodule", "/usr/lib/redis/modules/redisai.so"]
+CMD ["/usr/local/bin/redis-server", "--loadmodule", "/usr/lib/redis/modules/redisai.so"]

get_deps.sh

@@ -16,6 +16,7 @@ if [[ $1 == --help || $1 == help ]]; then
 		Argument variables:
 		CPU=1              Get CPU dependencies
 		GPU=1              Get GPU dependencies
+		JETSON=1           Get Jetson Nano dependencies
 		VERBOSE=1          Print commands
 		FORCE=1            Download even if present
 		WITH_DLPACK=0      Skip dlpack
@@ -106,6 +107,9 @@ if [[ $WITH_TF != 0 ]]; then
 				LIBTF_URL_BASE=https://storage.googleapis.com/tensorflow/libtensorflow
 			elif [[ $ARCH == arm64v8 ]]; then
 				TF_ARCH=arm64
+				if [[ $JETSON == 1 ]]; then
+					TF_BUILD="gpu-jetson"
+				fi
 				LIBTF_URL_BASE=https://s3.amazonaws.com/redismodules/tensorflow
 			elif [[ $ARCH == arm32v7 ]]; then
 				TF_ARCH=arm
@@ -205,7 +209,11 @@ if [[ $WITH_PT != 0 ]]; then
 			if [[ $GPU != 1 ]]; then
 				PT_BUILD=cpu
 			else
-				PT_BUILD=cu101
+				if [[ $JETSON == 1 ]]; then
+					PT_BUILD=cu102-jetson
+				else
+					PT_BUILD=cu101
+				fi
 			fi
 			if [[ $ARCH == x64 ]]; then
 				PT_ARCH=x86_64

opt/Makefile

@@ -53,6 +53,7 @@ make flow_tests  # run flow test
   COV=1            # perform coverage analysis
   VALGRIND|VGD=1   # test with Valgrind (implies DEBUG=1)
   CALLGRIND|CGD=1  # test with Callgrind (implies DEBUG=1)
+  REDIS=addr       # Address of Redis server for testing
 
 make cov-upload   # upload coverage data to codecov.io (requires CODECOV_TOKEN)
 
@@ -75,6 +76,7 @@ device selection options (fetch, build, and test):
   CPU=1         # build for CPU
   GPU=1         # build for GPU
   CUDA=1        # build for GPU
+  JETSON=1      # build for Nvidia Jetson
 endef
 
 #---------------------------------------------------------------------------------------------- 
@@ -85,10 +87,13 @@ ifeq ($(GPU),1)
 ifeq ($(CPU),1)
 $(error CPU=1 and GPU=1 (or CUDA=1) are conflicting)
 endif
-DEPS_FLAGS=gpu
+DEPS_FLAGS += GPU=1
 DEVICE=gpu
+ifeq ($(JETSON),1)
+DEPS_FLAGS += JETSON=1
+endif
 else
-DEPS_FLAGS=cpu
+DEPS_FLAGS += CPU=1
 DEVICE=cpu
 endif
 
@@ -194,7 +199,7 @@ setup:
 
 fetch deps:
 	@echo Fetching dependencies...
-	$(SHOW)VERBOSE=$(_SHOW) $(ROOT)/get_deps.sh $(DEPS_FLAGS)
+	$(SHOW)VERBOSE=$(_SHOW) $(DEPS_FLAGS) $(ROOT)/get_deps.sh
 
 format:
 	./clang-format-all.sh
@@ -243,6 +248,7 @@ flow_tests: build
 		CLUSTER=$(CLUSTER) \
 		GEN=$(GEN) AOF=$(AOF) SLAVES=$(SLAVES) \
 		VALGRIND=$(VALGRIND) \
+		REDIS=$(REDIS) \
 		$(ROOT)/tests/flow/tests.sh
 	$(COVERAGE_COLLECT_REPORT)
 
@@ -257,6 +263,7 @@ test: build
 		CLUSTER=$(CLUSTER) \
 		GEN=$(GEN) AOF=$(AOF) SLAVES=$(SLAVES) \
 		VALGRIND=$(VALGRIND) \
+		REDIS=$(REDIS) \
 		$(ROOT)/tests/flow/tests.sh
 	$(COVERAGE_COLLECT_REPORT)
 

opt/build/docker/Makefile

@@ -9,7 +9,7 @@ OSNICK ?= buster
 REPO=redisfab
 STEM=$(REPO)/redisai
 
-REDIS_VER:=6.0.5
+REDIS_VER:=6.0.9
 
 ART_DIR=$(ROOT)/bin/artifacts
 ART_INT_DIR=/var/opt/redislabs/artifacts
@@ -23,6 +23,9 @@ BUILD_ARGS.arm32v7=WITH_ORT=0 WITH_PT=0
 ifeq ($(CPU),1)
 VAR=cpu
 
+FETCH_ARGS=$(FETCH_ARGS.$(ARCH))
+BUILD_ARGS=$(BUILD_ARGS.$(ARCH))
+
 else ifeq ($(GPU),1)
 VAR=gpu
 DOCKERFILE_STEM ?= $(ROOT)/Dockerfile.gpu
@@ -33,12 +36,20 @@ OS.centos7=centos7
 OS.bionic=ubuntu18.04
 OS.xenial=ubuntu16.04
 
-else
-VAR=cpu
+FETCH_ARGS=$(FETCH_ARGS.$(ARCH))
+BUILD_ARGS=$(BUILD_ARGS.$(ARCH))
+
+else ifeq ($(JETSON),1)
+VAR=jetson
+DOCKERFILE_STEM ?= $(ROOT)/Dockerfile.jetson
+
+# FETCH_ARGS=WITH_TF=1 WITH_PT=1 WITH_TFLITE=0 WITH_ORT=0
+# BUILD_ARGS=WITH_TF=1 WITH_PT=1 WITH_TFLITE=0 WITH_ORT=0
+
 endif
 
-DOCKER_BUILD_ARGS += \
-	DEPS_ARGS="$(FETCH_ARGS.$(ARCH))" \
-	BUILD_ARGS="$(BUILD_ARGS.$(ARCH))"
+# DOCKER_BUILD_ARGS += \
+#	DEPS_ARGS="$(FETCH_ARGS)" \
+#	BUILD_ARGS="$(BUILD_ARGS)"
 
 include $(MK)/docker

opt/build/libtorch/Makefile

@@ -1,91 +1,120 @@
 
 ROOT=../../..
 
-VERSION ?= 1.5.0
-OSNICK ?= buster
+VERSION ?= 1.7.0
+OSNICK ?= bionic
+
+REDIS_VER=6.0.9
 
 #----------------------------------------------------------------------------------------------
 
 S3_URL=redismodules/pytorch
 
 OS:=$(shell $(ROOT)/opt/readies/bin/platform --os)
-ifeq ($(OS),macosx)
-OS:=macos
-endif
-STEM=libtorch-cpu-$(OS)
+OSNICK:=$(shell $(ROOT)/opt/readies/bin/platform --osnick)
+ARCH:=$(shell $(ROOT)/opt/readies/bin/platform --arch)
+DEVICE ?= cpu
+
+STEM=libtorch
 
 DOCKER_OS.bionic=ubuntu:bionic
 DOCKER_OS.stretch=debian:stretch-slim
 DOCKER_OS.buster=debian:buster-slim
 DOCKER_OS=$(DOCKER_OS.$(OSNICK))
 
+TAR_ARCH.x64=x86_64
+TAR_ARCH.arm64v8=arm64
+TAR_ARCH.arm32v7=arm
+
 #----------------------------------------------------------------------------------------------
 
 define targets # (1=OP, 2=op)
 $(1)_TARGETS :=
-$(1)_TARGETS += $(if $(findstring $(X64),1),$(2)_x64)
+$(1)_TARGETS += $(if $(findstring $(X64),1),$(2)_native)
+ifeq ($$(CROSS),1)
 $(1)_TARGETS += $(if $(findstring $(ARM7),1),$(2)_arm32v7)
 $(1)_TARGETS += $(if $(findstring $(ARM8),1),$(2)_arm64v8)
+endif
 
-$(1)_TARGETS += $$(if $$(strip $$($(1)_TARGETS)),,$(2)_x64 $(2)_arm32v7 $(2)_arm64v8)
+ifeq ($$(CROSS),1)
+$(1)_TARGETS += $$(if $$(strip $$($(1)_TARGETS)),,$(2)_native $(2)_arm32v7 $(2)_arm64v8)
+else
+$(1)_TARGETS += $$(if $$(strip $$($(1)_TARGETS)),,$(2)_native)
+endif
 endef
 
 $(eval $(call targets,BUILD,build))
 $(eval $(call targets,PUBLISH,publish))
 
 #----------------------------------------------------------------------------------------------
 
-define build_x64 # (1=arch, 2=tar-arch)
+DOCKER_BUILD_ARGS += \
+	PT_VER=$(VERSION) \
+	OSNICK=$(OSNICK) \
+	OS=$(DOCKER_OS) \
+	ARCH=$(ARCH)
+
+define build_native # (1=arch, 2=tar-arch)
 IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
 
-build_x64:
-	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64 \
-		--build-arg OS=$(DOCKER_OS) $(ROOT)
+build_native:
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM):$(VERSION)-$(DEVICE)-$(1) -f Dockerfile.$(VARIANT) \
+		$(foreach A,$(DOCKER_BUILD_ARGS),--build-arg $(A)) .
 	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
-	@docker cp `cat $$(CID_$(1))`:/build/dest/$(STEM)-$(2)-$(VERSION).tar.gz .
+	@docker cp `cat $$(CID_$(1))`:/build/dest/$(STEM)-$(DEVICE)-$(OS)-$(2)-$(VERSION).tar.gz .
 
-.PHONY: build_x64
+.PHONY: build_native
 endef
 
+ifeq ($(CROSS),1)
+
 define build_arm # (1=arch, 2=tar-arch)
 IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
 
 build_$(1):
-	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.arm \
-		--build-arg ARCH=$(1) $(ROOT)
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM):$(VERSION)-$(DEVICE)-$(1) -f Dockerfile.$(VARIANT) \
+		$(foreach A,$(DOCKER_BUILD_ARGS),--build-arg $(A)) .
 	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
 	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(2)-$(VERSION).tar.gz .
 
 .PHONY: build_$(1)
 endef
 
+endif # CROSS
+
 #----------------------------------------------------------------------------------------------
 
-define publish_x64 # (1=arch, 2=tar-arch)
-publish_x64:
+define publish_native # (1=arch, 2=tar-arch)
+publish_native:
 	@aws s3 cp $(STEM)-$(2)-$(VERSION).tar.gz s3://$(S3_URL)/ --acl public-read
 
-.PHONY: publish_x64
+.PHONY: publish_native
 endef
 
+ifeq ($(CROSS),1)
+
 define publish_arm # (1=arch, 2=tar-arch)
 publish_$(1):
 	@aws s3 cp $(STEM)-$(2)-$(VERSION).tar.gz s3://$(S3_URL)/ --acl public-read
 
 .PHONY: publish_$(1)
 endef
 
+endif # CROSS
+
 #----------------------------------------------------------------------------------------------
 
 all: build publish
 
 build: $(BUILD_TARGETS)
 
-$(eval $(call build_x64,x64,x86_64))
+$(eval $(call build_native,$(ARCH),$(TAR_ARCH.$(ARCH))))
+ifeq ($(CROSS),1)
 $(eval $(call build_arm,arm64v8,arm64))
 $(eval $(call build_arm,arm32v7,arm))
+endif
 
 ifneq ($(filter publish,$(MAKECMDGOALS)),)
 ifeq ($(wildcard $(HOME)/.aws/credentials),)
@@ -96,13 +125,16 @@ endif
 publish: $(PUBLISH_TARGETS)
 
 $(eval $(call publish_x64,x64,x86_64))
+ifeq ($(CROSS),1)
 $(eval $(call publish_arm,arm64v8,arm64))
 $(eval $(call publish_arm,arm32v7,arm))
+endif
 
 repack:
 	@PT_VERSION=$(VERSION) GPU=$(GPU) ./repack.sh
 
 help:
-	@echo "make [build|repack|publish] [X64=1|ARM7=1|ARM8=1]"
+	@echo "make [build|repack|publish] [X64=1|ARM7=1|ARM8=1] [CROSS=1]"
 
 .PHONY: all build repack publish help
+