Add environment variables to DALI TF build image (NVIDIA#1034)

Signed-off-by: Joaquin Anton <[email protected]>
zhwzhong · Jul 3, 2019 · 257650b · 257650b
1 parent 8ceabdb
commit 257650b
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/dali_tf_plugin/CMakeLists.txt b/dali_tf_plugin/CMakeLists.txt
@@ -24,6 +24,9 @@ if (DALI_BUILD_FLAVOR)
   set(DALI_VERSION "${DALI_VERSION}.${TIMESTAMP}")
 endif()
 
+# Note: CUDA 8 support is unofficial. CUDA 9 is officially supported
+find_package(CUDA 8.0 REQUIRED)
+
 configure_file("${PROJECT_SOURCE_DIR}/setup.py.in" "${PROJECT_BINARY_DIR}/setup.py")
 file(GLOB PREBUILT_PLUGINS "libdali_tf_*.so")
 file(COPY ${PREBUILT_PLUGINS} DESTINATION "${PROJECT_BINARY_DIR}")

diff --git a/dali_tf_plugin/build_in_custom_op_docker.sh b/dali_tf_plugin/build_in_custom_op_docker.sh
@@ -33,7 +33,12 @@ done
 
 mkdir -p dali_tf_sdist_build
 cd dali_tf_sdist_build
-cmake ..
+
+cmake .. \
+      -DDALI_BUILD_FLAVOR=${NVIDIA_DALI_BUILD_FLAVOR} \
+      -DTIMESTAMP=${DALI_TIMESTAMP} \
+      -DGIT_SHA=${GIT_SHA}
+
 make -j
 python setup.py sdist
 cp dist/*.tar.gz /dali_tf_sdist
diff --git a/docker/Dockerfile_dali_tf b/docker/Dockerfile_dali_tf
@@ -3,6 +3,15 @@ FROM ${TF_CUSTOM_OP_BUILDER_IMAGE}
 
 WORKDIR /opt/dali_tf_build
 
+ARG NVIDIA_DALI_BUILD_FLAVOR
+ENV NVIDIA_DALI_BUILD_FLAVOR=${NVIDIA_DALI_BUILD_FLAVOR}
+ARG GIT_SHA
+ENV GIT_SHA=${GIT_SHA}
+ARG DALI_TIMESTAMP
+ENV DALI_TIMESTAMP=${DALI_TIMESTAMP}
+ARG NVIDIA_BUILD_ID
+ENV NVIDIA_BUILD_ID ${NVIDIA_BUILD_ID:-0}
+
 RUN mkdir /dali_tf_sdist && chmod 0777 /dali_tf_sdist
 
 COPY Acknowledgements.txt .

diff --git a/docker/build.sh b/docker/build.sh
@@ -173,7 +173,14 @@ export CUSTOM_OP_BUILDER_CONTAINER="${CUSTOM_OP_BUILDER_IMAGE_NAME}_container"
 
 mkdir -p dali_tf_plugin/whl
 cp ${tmp_wheelhouse}/*.whl dali_tf_plugin/whl/
-docker build -t ${CUSTOM_OP_BUILDER_IMAGE_NAME} -f docker/Dockerfile_dali_tf --build-arg "TF_CUSTOM_OP_BUILDER_IMAGE=${CUSTOM_OP_BUILDER_CLEAN_IMAGE_NAME}"  .
+docker build -t ${CUSTOM_OP_BUILDER_IMAGE_NAME} -f docker/Dockerfile_dali_tf \
+             --build-arg "TF_CUSTOM_OP_BUILDER_IMAGE=${CUSTOM_OP_BUILDER_CLEAN_IMAGE_NAME}" \
+             --build-arg "NVIDIA_BUILD_ID=${NVIDIA_BUILD_ID}" \
+             --build-arg "NVIDIA_DALI_BUILD_FLAVOR=${DALI_BUILD_FLAVOR}" \
+             --build-arg "GIT_SHA=${GIT_SHA}" \
+             --build-arg "DALI_TIMESTAMP=${DALI_TIMESTAMP}" \
+             .
+
 nvidia-docker run --name ${CUSTOM_OP_BUILDER_CONTAINER} ${CUSTOM_OP_BUILDER_IMAGE_NAME}
 tmp_dali_tf_sdist=$(mktemp -d)
 docker cp "${CUSTOM_OP_BUILDER_CONTAINER}:/dali_tf_sdist/" "${tmp_dali_tf_sdist}"