Security-Tools-Alliance · psyray · Nov 9, 2024 · Nov 9, 2024 · Nov 10, 2024 · Nov 10, 2024
diff --git a/.env-dist b/.env-dist
@@ -43,3 +43,13 @@ INSTALL_TYPE=prebuilt
 DJANGO_SUPERUSER_USERNAME=rengine
 DJANGO_SUPERUSER_EMAIL=rengine@example.com
 DJANGO_SUPERUSER_PASSWORD=Sm7IJG.IfHAFw9snSKv
+
+#
+# GPU Configuration for Ollama LLM
+# Set GPU=1 to enable GPU support (default: 0 for CPU only)
+# GPU_TYPE will be automatically set during installation (nvidia or rocm)
+# DOCKER_RUNTIME will match GPU_TYPE when GPU support is enabled
+#
+GPU=0
+GPU_TYPE=none
+DOCKER_RUNTIME=none
diff --git a/Makefile b/Makefile
@@ -14,6 +14,7 @@ COMPOSE_FILE	      := docker/docker-compose.yml
 COMPOSE_FILE_BUILD	  := docker/docker-compose.build.yml
 COMPOSE_FILE_DEV      := docker/docker-compose.dev.yml
 COMPOSE_FILE_SETUP    := docker/docker-compose.setup.yml
+COMPOSE_FILE_GPU      := docker/docker-compose.gpu.yml
 SERVICES              := db web proxy redis celery celery-beat ollama
 
 # Check if 'docker compose' command is available, otherwise check for 'docker-compose'
@@ -36,7 +37,33 @@ $(info Using: $(DOCKER_COMPOSE))
 DOCKER_COMPOSE_CMD      := ${COMPOSE_PREFIX_CMD} ${DOCKER_COMPOSE}
 DOCKER_COMPOSE_FILE_CMD := ${DOCKER_COMPOSE_CMD} -f ${COMPOSE_FILE}
 
-# --------------------------
+# Add GPU variable with default value
+GPU ?= 0
+
+# Function to handle GPU configuration
+define gpu_config
+	$(info Checking GPU configuration...)
+	$(if $(filter 1,$(GPU)), \
+		$(info GPU=1, detecting GPU type...) \
+		$(eval GPU_TYPE := $(shell ./scripts/gpu_support.sh)) \
+		$(if $(filter nvidia,$(GPU_TYPE)), \
+			$(info Configuring for NVIDIA GPU) \
+			$(eval DOCKER_RUNTIME := nvidia) \
+			$(eval COMPOSE_GPU_FILE := -f ${COMPOSE_FILE_GPU} --profile gpu), \
+			$(if $(filter amd,$(GPU_TYPE)), \
+				$(info Configuring for AMD GPU) \
+				$(eval DOCKER_RUNTIME := amd) \
+				$(eval COMPOSE_GPU_FILE := -f ${COMPOSE_FILE_GPU} --profile gpu), \
+				$(info No supported GPU detected) \
+				$(eval COMPOSE_GPU_FILE :=) \
+			) \
+		), \
+		$(info GPU support disabled) \
+		$(eval COMPOSE_GPU_FILE :=) \
+	)
+	$(eval export GPU_TYPE) \
+	$(eval export DOCKER_RUNTIME)
+endef
 
 .PHONY: certs up dev_up build_up build pull superuser_create superuser_delete superuser_changepassword migrate down stop restart remove_images test logs images prune help
 
@@ -46,9 +73,10 @@ pull:			## Pull pre-built Docker images from repository.
 images:			## Show all Docker images for reNgine services.
 	@docker images --filter=reference='ghcr.io/security-tools-alliance/rengine-ng:*' --format "table {{.Repository}}\t{{.Tag}}\t{{.ID}}\t{{.Size}}"
 
-build:			## Build all Docker images locally.
+build:			## Build all Docker images locally. Use GPU=1 to enable GPU support.
 	@make remove_images
-	${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_BUILD} build ${SERVICES}
+	$(call gpu_config)
+	${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_BUILD} ${COMPOSE_GPU_FILE} build ${SERVICES}
 
 build_up:		## Build and start all services.
 	@make down
@@ -58,12 +86,14 @@ build_up:		## Build and start all services.
 certs:		    ## Generate certificates.
 	@${DOCKER_COMPOSE_CMD} -f ${COMPOSE_FILE_SETUP} run --rm certs
 
-up:				## Pull and start all services.
-	${DOCKER_COMPOSE_FILE_CMD} up -d ${SERVICES}
+up:				## Pull and start all services. Use GPU=1 to enable GPU support.
+	$(call gpu_config)
+	${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} up -d ${SERVICES}
 
-dev_up:			## Pull and start all services with development configuration (more debug logs and Django Toolbar in UI).
+dev_up:			## Pull and start all services with development configuration. Use GPU=1 to enable GPU support.
 	@make down
-	${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} up -d ${SERVICES}
+	$(call gpu_config)
+	${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} up -d ${SERVICES}
 
 superuser_create:		## Generate username (use only after `make up`).
 ifeq ($(isNonInteractive), true)
@@ -92,44 +122,45 @@ stop:			## Stop all services.
 	${DOCKER_COMPOSE_FILE_CMD} stop ${SERVICES}
 
 restart:		## Restart specified services or all if not specified. Use DEV=1 for development mode, COLD=1 for down and up instead of restart.
+	$(call gpu_config)
 	@if [ "$(COLD)" = "1" ]; then \
 		if [ "$(DEV)" = "1" ]; then \
 			if [ -n "$(filter-out $@,$(MAKECMDGOALS))" ]; then \
 				echo "Cold restart $(filter-out $@,$(MAKECMDGOALS)) in dev mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} down $(filter-out $@,$(MAKECMDGOALS)); \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} up -d $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} down $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} up -d $(filter-out $@,$(MAKECMDGOALS)); \
 			else \
 				echo "Cold restart ${SERVICES} in dev mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} down; \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} up -d ${SERVICES}; \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} down; \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} up -d ${SERVICES}; \
 			fi \
 		else \
 			if [ -n "$(filter-out $@,$(MAKECMDGOALS))" ]; then \
 				echo "Cold restart $(filter-out $@,$(MAKECMDGOALS)) in production mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} down $(filter-out $@,$(MAKECMDGOALS)); \
-				${DOCKER_COMPOSE_FILE_CMD} up -d $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} down $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} up -d $(filter-out $@,$(MAKECMDGOALS)); \
 			else \
 				echo "Cold restart ${SERVICES} in production mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} down; \
-				${DOCKER_COMPOSE_FILE_CMD} up -d ${SERVICES}; \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} down; \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} up -d ${SERVICES}; \
 			fi \
 		fi \
 	else \
 		if [ "$(DEV)" = "1" ]; then \
 			if [ -n "$(filter-out $@,$(MAKECMDGOALS))" ]; then \
 				echo "Restart $(filter-out $@,$(MAKECMDGOALS)) in dev mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} restart $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} restart $(filter-out $@,$(MAKECMDGOALS)); \
 			else \
 				echo "Restart ${SERVICES} in dev mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} restart ${SERVICES}; \
+				${DOCKER_COMPOSE_FILE_CMD} -f ${COMPOSE_FILE_DEV} ${COMPOSE_GPU_FILE} restart ${SERVICES}; \
 			fi \
 		else \
 			if [ -n "$(filter-out $@,$(MAKECMDGOALS))" ]; then \
 				echo "Restart $(filter-out $@,$(MAKECMDGOALS)) in production mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} restart $(filter-out $@,$(MAKECMDGOALS)); \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} restart $(filter-out $@,$(MAKECMDGOALS)); \
 			else \
 				echo "Restart ${SERVICES} in production mode"; \
-				${DOCKER_COMPOSE_FILE_CMD} restart ${SERVICES}; \
+				${DOCKER_COMPOSE_FILE_CMD} ${COMPOSE_GPU_FILE} restart ${SERVICES}; \
 			fi \
 		fi \
 	fi
@@ -159,7 +190,10 @@ help:			## Show this help.
 	@echo "Manage Docker images, containers and Django commands using Docker Compose files."
 	@echo ""
 	@echo "Usage:"
-	@echo "  make <target> (default: help)"
+	@echo "  make <target> [GPU=1] (default: help)"
+	@echo ""
+	@echo "Options:"
+	@echo "  GPU=1                                    Enable GPU support for Ollama LLM"
 	@echo ""
 	@echo "Targets:"
 	@echo "  make restart [service1] [service2] ...  				Restart specific services in production mode"
@@ -170,6 +204,12 @@ help:			## Show this help.
 	@echo "  make restart DEV=1 COLD=1 [service1] [service2] ...  	Cold restart (recreate containers) specific services in development mode"
 	@echo "  make restart COLD=1                     				Cold restart (recreate containers) all services in production mode"
 	@echo "  make restart DEV=1 COLD=1               				Cold restart (recreate containers) all services in development mode"
+	@echo ""
+	@echo "Examples:"
+	@echo "  make up GPU=1                          				Start all services with GPU support"
+	@echo "  make dev_up GPU=1                      				Start development environment with GPU support"
+	@echo "  make build GPU=1                       				Build all images with GPU support"
+	@echo "  make build_up GPU=1                    				Build and start all services with GPU support"
 
 %:
 	@:
diff --git a/docker/docker-compose.gpu.yml b/docker/docker-compose.gpu.yml
@@ -0,0 +1,21 @@
+services:
+  ollama:
+    extends:
+      file: docker-compose.yml
+      service: ollama
+    build:
+      context: ./ollama
+      dockerfile: ${GPU_TYPE:+${GPU_TYPE}.Dockerfile}
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    environment:
+      # NVIDIA Settings
+      - NVIDIA_VISIBLE_DEVICES=all
+      - NVIDIA_DRIVER_CAPABILITIES=all
+    profiles:
+      - gpu
diff --git a/docker/ollama/Dockerfile b/docker/ollama/Dockerfile
@@ -1 +1,5 @@
-FROM ollama/ollama:0.3.6
+FROM ollama/ollama:0.3.6
+
+# Add labels for GPU support information
+LABEL com.rengine.gpu-support="optional"
+LABEL com.rengine.gpu-description="GPU support can be enabled using docker-compose.gpu.yml"
diff --git a/docker/ollama/amd.Dockerfile b/docker/ollama/amd.Dockerfile
@@ -0,0 +1,9 @@
+FROM ollama/ollama:0.3.6
+
+# Enable ROCm support
+ENV HSA_OVERRIDE_GFX_VERSION=9.0.0
+ENV ROCR_VISIBLE_DEVICES=all
+
+# Add labels for GPU support information
+LABEL com.rengine.gpu-support="enabled-amd"
+LABEL com.rengine.gpu-description="AMD GPU support is enabled in this image"
diff --git a/docker/ollama/nvidia.Dockerfile b/docker/ollama/nvidia.Dockerfile
@@ -0,0 +1,9 @@
+FROM ollama/ollama:0.3.6
+
+# Enable NVIDIA GPU support
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+# Add labels for GPU support information
+LABEL com.rengine.gpu-support="enabled-nvidia"
+LABEL com.rengine.gpu-description="NVIDIA GPU support is enabled in this image"
diff --git a/docker/web/entrypoint.sh b/docker/web/entrypoint.sh
@@ -4,6 +4,11 @@
 poetry run -C $HOME/ python3 manage.py collectstatic --noinput
 
 # Run production server
-poetry run -C $HOME/ gunicorn reNgine.wsgi:application -w 8 --bind 0.0.0.0:8000 --limit-request-line 0
+poetry run -C $HOME/ gunicorn reNgine.wsgi:application \
+    --workers 8 \
+    --bind 0.0.0.0:8000 \
+    --limit-request-line 0 \
+    --timeout 120 \
+    --keep-alive 120
 
 exec "$@"