abc123 / Dockerfile.jetson
vimalk78's picture
fix: add CUDA warmup and memory config for Jetson GPU support
4a0fccf
Raw
History Blame Contribute Delete
2.72 kB
# Dockerfile for NVIDIA Jetson (ARM64) - Jetson Orin Nano, Xavier, etc.
# Uses NVIDIA's L4T PyTorch container as base for proper GPU support
#
# Check your JetPack/L4T version: cat /etc/nv_tegra_release
#
# Available base images:
# JetPack 6.x (L4T R36.x): dustynv/l4t-pytorch:r36.4.0
# JetPack 5.1.x (L4T R35.x): nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
FROM dustynv/l4t-pytorch:r36.4.0
WORKDIR /app
# Install Node.js for frontend build
RUN apt-get update && apt-get install -y \
curl \
git \
&& curl -fsSL https://deb.nodesource.com/setup_18.x | bash - \
&& apt-get install -y nodejs \
&& rm -rf /var/lib/apt/lists/*
# Copy and build frontend
COPY crossword-app/frontend/package*.json ./frontend/
RUN cd frontend && npm ci
COPY crossword-app/frontend/ ./frontend/
RUN cd frontend && npm run build
# Install Python dependencies (PyTorch already in base image)
# Use Jetson-specific requirements with Python 3.10 compatible versions
COPY crossword-app/backend-py/requirements-jetson.txt ./backend-py/
RUN pip3 install --no-cache-dir --index-url https://pypi.org/simple -r backend-py/requirements-jetson.txt && \
pip3 install --no-cache-dir --no-deps --index-url https://pypi.org/simple sentence-transformers>=2.2.0
# Copy backend code
COPY crossword-app/backend-py/ ./backend-py/
COPY crossword-app/words/ ./backend-py/words/
# Copy cache directory with pre-built models and NLTK data
# Exclude embeddings .pt file - it was created on x86 and has CUDA compatibility issues
# Jetson will regenerate embeddings on first run (takes ~2-3 minutes)
COPY cache-dir/models--sentence-transformers--all-mpnet-base-v2/ ./backend-py/cache/models--sentence-transformers--all-mpnet-base-v2/
COPY cache-dir/nltk_data/ ./backend-py/cache/nltk_data/
COPY cache-dir/norvig_*.pkl ./backend-py/cache/
RUN chmod -R 755 ./backend-py/cache/ || true
# Copy built frontend to backend public directory
RUN mkdir -p backend-py/public && cp -r frontend/dist/* backend-py/public/
WORKDIR /app/backend-py
EXPOSE 7860
ENV NODE_ENV=production
ENV PORT=7860
ENV PYTHONPATH=/app/backend-py
ENV PYTHONUNBUFFERED=1
ENV CACHE_DIR=/app/backend-py/cache
ENV NLTK_DATA=/app/backend-py/cache/nltk_data
ENV VOCAB_SOURCE=norvig
ENV NORVIG_VOCAB_PATH=/app/backend-py/words/norvig/count_1w100k.txt
# CUDA memory allocation config for Jetson unified memory
ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# Model: all-mpnet-base-v2 (420MB, best quality) or all-MiniLM-L6-v2 (90MB, faster)
# Set THEMATIC_MODEL_NAME=all-MiniLM-L6-v2 if you encounter GPU memory issues
ENV THEMATIC_MODEL_NAME=all-mpnet-base-v2
CMD ["python3", "-m", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]