From f0fce2dc54157f63d476fd123f5fc41048b00a08 Mon Sep 17 00:00:00 2001
From: Alex Davies <traverse.da@gmail.com>
Date: Thu, 19 Sep 2024 11:25:59 -0300
Subject: [PATCH] Added flag for enabling nvidia gpu

---
 copier.yml                   | 5 +++++
 src/docker-compose.yml.jinja | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/copier.yml b/copier.yml
index 48640b6..7bd2481 100644
--- a/copier.yml
+++ b/copier.yml
@@ -24,6 +24,11 @@ LANGUAGES:
 
 RUN_COMMAND: rosrun {{ PACKAGE_NAME }} {{ "main.py" if 'Python' in LANGUAGES else "main" if "CPP" in LANGUAGES else "yourcommandGoes here" }}
 
+ENABLE_NVIDIA_GPU:
+  default: False
+  type: bool
+  help: Enabling nvidia gpus make deployment much more complicated. Ensure you have the nvidia container toolkit working properly.
+
 HEALTHCHECK:
   defualt: --start-period=60s --start-interval=1s CMD /ros_entrypoint.sh rostopic list
   type: str
diff --git a/src/docker-compose.yml.jinja b/src/docker-compose.yml.jinja
index 42e318b..3b474b2 100644
--- a/src/docker-compose.yml.jinja
+++ b/src/docker-compose.yml.jinja
@@ -19,6 +19,15 @@ services:
       nofile:
         soft: 1024
         hard: 524288
+{% if ENABLE_GPU %}
+    deploy:
+      resources:
+        reservations:
+          devices:
+           - driver: cdi
+             device_ids:
+                - nvidia.com/gpu=all
+{{% endif %}
 
 networks:
   sdk: