Multi-stage build in Dockerfile, reduces size by 3GB, and runs tests

Dockerfile

@@ -1,8 +1,14 @@
-FROM python:3.10-slim
+FROM python:3.10-slim as base
 
 WORKDIR /comma-fixer
+ENV PYTHONUNBUFFERED=1
+
+RUN python -m venv /venv
+ENV PATH="/venv/bin:$PATH"
 
 COPY requirements.txt .
+COPY test-requirements.txt .
+RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
 COPY src/baseline.py src/baseline.py
@@ -10,5 +16,20 @@ RUN python src/baseline.py  # This pre-downloads models and tokenizers
 
 COPY . .
 
+FROM base as test
+
+RUN pip install -r test-requirements.txt
+WORKDIR src
+RUN python -m pytest ../tests
+
+FROM python:3.10-slim as deploy
+
+WORKDIR /comma-fixer
+COPY --from=base /comma-fixer /comma-fixer
+COPY --from=base /venv /venv
+# Copy pre-downloaded models and make sure we are using the env
+COPY --from=base /root/.cache/huggingface/hub/ /root/.cache/huggingface/hub/
+ENV PATH="/venv/bin:$PATH"
+
 EXPOSE 8000
 CMD uvicorn "app:app" --port 8000 --host "0.0.0.0"

README.txt

@@ -1 +1,3 @@
-`docker log [id]` for logs from the container.
+`docker log [id]` for logs from the container.
+
+`docker build -t comma-fixer --target test .` for tests

docker-compose.yml

@@ -19,8 +19,10 @@ services:
       dockerfile: Dockerfile
     container_name: comma-fixer
     command: uvicorn --host 0.0.0.0 --port 8000 "app:app"
-    volumes:
-      - ./:/comma-fixer
+    ports:
+      - "8000:8000"
+#    volumes:
+#      - ./:/comma-fixer
 #    networks:
 #      my-network:
 #        aliases:

requirements.txt

@@ -1,8 +1,6 @@
 fastapi==0.101.1
 gunicorn==21.2.0
 uvicorn==0.23.2
-pytest
-httpx
 torch==2.0.1
 transformers==4.31.0
 

src/baseline.py

@@ -33,7 +33,7 @@ def _fix_commas_based_on_pipeline_output(pipeline_json: list[dict], original_s:
         current_offset = _find_current_token(current_offset, i, pipeline_json, result)
         if _should_insert_comma(i, pipeline_json):
             result = result[:current_offset] + ',' + result[current_offset:]
-        current_offset += 1
+            current_offset += 1
     return result
 
 

test-requirements.txt

@@ -0,0 +1,2 @@
+pytest
+httpx