Dockerfile.gpu

@@ -0,0 +1,11 @@
+FROM pytorch/pytorch:1.7.1-cuda11.0-cudnn8-runtime
+
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+ADD . /classla-api
+WORKDIR /classla-api
+
+RUN pip install $(grep -ivE "torch" requirements.txt)
+
+CMD [ "python3", "-m" , "flask", "run", "--host=0.0.0.0"]

docker-compose.yml

@@ -1,15 +1,13 @@
-version: "3.8"
+version: "3.7"
 services:
   app:
     build: .
     command: python -m flask run --host=0.0.0.0
-    deploy:
-      resources:
-        reservations:
-          devices:
-            - driver: nvidia
-              count: all
-              capabilities: [ gpu ]
+    runtime: nvidia
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=all
+      - NVIDIA_DRIVER_CAPABILITIES=all
+      - CUDA_VISIBLE_DEVICES=1
     ports:
       - "5000:5000"
     volumes: