Skip to content

Commit 6a98d7f

Browse files
committed
Fix CuPy JIT: use CUDA devel image for runtime stage
CuPy needs cuda_fp16.h for JIT kernel compilation at runtime. The runtime CUDA image doesn't include development headers. Also adds --no-tar flag to pipeline arg parser.
1 parent 98277c8 commit 6a98d7f

2 files changed

Lines changed: 4 additions & 1 deletion

File tree

docker/Dockerfile

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,9 @@ RUN git clone https://github.com/tomography/tomocupy.git /tmp/tomocupy \
4848

4949

5050
# -- Stage 2: Runtime --
51-
FROM nvidia/cuda:12.2.2-runtime-ubuntu22.04
51+
# Using devel (not runtime) because CuPy JIT-compiles CUDA kernels at runtime
52+
# and needs headers like cuda_fp16.h from the CUDA toolkit
53+
FROM nvidia/cuda:12.2.2-devel-ubuntu22.04
5254

5355
ENV DEBIAN_FRONTEND=noninteractive
5456

pipeline/run_full_pipeline.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ while [[ $# -gt 0 ]]; do
9494
--sam3-stability-score-thresh) SAM3_STABILITY_SCORE_THRESH="$2"; shift 2 ;;
9595
--sam3-min-mask-area) SAM3_MIN_MASK_AREA="$2"; shift 2 ;;
9696
--sam3-preview-slices) SAM3_PREVIEW_SLICES="$2"; shift 2 ;;
97+
--no-tar) shift ;; # accepted but no-op on AWS (no tar step)
9798
*) echo "WARNING: Unknown argument $1"; shift ;;
9899
esac
99100
done

0 commit comments

Comments
 (0)