feat: build_ubuntu.sh update

build_space/build_ubuntu.sh

-#/bin/bash
-
-# pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=12.1
-# cuda=11.6（这个版本的空镜像，有些框架需要低版本）
-# paddlepaddle-gpu==2.6.1.post120  cuda12
-# tensroflow 安装最高版本 cuda12
-
-rm -r ./tmp
+#!/bin/bash
 
+# 框架
 framework=$1
-newtag=$2
+# 输出镜像tag
+image_tag=$2
+# 基础镜像
+base_image=$3
+build_args=" --build-arg BASE_IMAGE=$base_image"
+if [ ${base_image%%:*} = "pytorch/pytorch" ]; then
+    build_args="$build_args --build-arg BASE_IMAGE_IS_TORCH=1 "
+fi
+for arg in ${*:4}
+do
+    build_args="$build_args --build-arg $arg "
+done
+
 tmp_dockerfile="Dockerfile.${RANDOM}"
 
 if [ ! -d "./tmp" ];then
@@ -28,10 +34,14 @@ cp ./python-requirements.txt ./tmp/
 cp -f ./Dockerfile.${framework}_ubuntu ./tmp/${tmp_dockerfile}
 # temp_image="$(grep -n '^FROM' ./tmp/${tmp_dockerfile} | tac | head -1 | awk '{print $2}')"
 # sed -i "s?${temp_image}?${base_tag}?g" ./tmp/${tmp_dockerfile}
-docker build -f ./tmp/${tmp_dockerfile} -t ${newtag} ./tmp/
+# docker build --build-arg BASE_IMAGE_IS_TORCH=1 --build-arg IMAGE="pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime" -f ./tmp/${tmp_dockerfile} -t ${newtag} ./tmp/
+echo "docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/"
+docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/
+#echo "docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/"
+rm -r ./tmp
 if [[ $? -eq 0 ]];then
         echo -e "\033[32mBuild Image Successfully !\033[0m"
 else
-        echo -e "\033[32mBuild Image unsuccessful !\033[0m"
+        echo -e "\033[32mBuild Image fail!\033[0m"
         exit 1
 fi
\ No newline at end of file