dl_checkpoints.sh run as root but chown to current user

livepeer · Jan 23, 2025 · a262b28 · a262b28
1 parent ad11b40
commit a262b28
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/runner/dl_checkpoints.sh b/runner/dl_checkpoints.sh
@@ -135,26 +135,29 @@ function build_tensorrt_models() {
     docker pull $AI_RUNNER_STREAMDIFFUSION_IMAGE
     # ai-worker has tags hardcoded in `var livePipelineToImage` so we need to use the same tag in here:
     docker image tag $AI_RUNNER_STREAMDIFFUSION_IMAGE livepeer/ai-runner:live-app-streamdiffusion
-    docker run --user $(id -u):$(id -g) --rm -v ./models:/models --gpus all -l TensorRT-engines \
-        $AI_RUNNER_STREAMDIFFUSION_IMAGE \
+    docker run --rm -v ./models:/models --gpus all -l TensorRT-engines $AI_RUNNER_STREAMDIFFUSION_IMAGE \
         bash -c "for model in $MODELS; do
                     for timestep in $TIMESTEPS; do
                         echo \"Building TensorRT engines for model=\$model timestep=\$timestep...\" && \
                         python app/live/StreamDiffusionWrapper/build_tensorrt.py --model-id \$model --timesteps \$timestep
                     done
-                done"
+                done
+                chown -R $(id -u -n):$(id -g -n) /models
+                " \
+        || (echo "failed streamdiffusion tensorrt"; return 1)
 
     # ComfyUI (only DepthAnything for now)
     AI_RUNNER_COMFYUI_IMAGE=${AI_RUNNER_COMFYUI_IMAGE:-livepeer/ai-runner:live-app-comfyui}
     docker pull $AI_RUNNER_COMFYUI_IMAGE
     # ai-worker has tags hardcoded in `var livePipelineToImage` so we need to use the same tag in here:
     docker image tag $AI_RUNNER_COMFYUI_IMAGE livepeer/ai-runner:live-app-comfyui
-    docker run --user $(id -u):$(id -g) --rm -v ./models:/models --gpus all -l TensorRT-engines \
-        $AI_RUNNER_COMFYUI_IMAGE \
+    docker run --rm -v ./models:/models --gpus all -l TensorRT-engines $AI_RUNNER_COMFYUI_IMAGE \
         bash -c "cd /comfyui/models/Depth-Anything-Onnx && \
                     python /comfyui/custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py && \
                     mkdir -p /comfyui/models/tensorrt/depth-anything && \
-                    mv *.engine /comfyui/models/tensorrt/depth-anything"
+                    mv *.engine /comfyui/models/tensorrt/depth-anything && \
+                    chown -R $(id -u -n):$(id -g -n) /models" \
+        || (echo "failed ComfyUI tensorrt"; return 1)
 }
 
 # Download models with a restrictive license.