From 7ba92b406b4490cbc890c182024e9d5bea0d6f1d Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Thu, 6 Feb 2025 16:57:33 +0100
Subject: [PATCH] fix: requesting gpu from docker runtime only for loaders that
 need it

---
 dyana/cli.py                         | 4 ++--
 dyana/loaders/automodel/settings.yml | 1 +
 dyana/loaders/loader.py              | 2 +-
 dyana/loaders/lora/settings.yml      | 2 ++
 dyana/loaders/ollama/settings.yml    | 2 ++
 dyana/loaders/settings.py            | 1 +
 6 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/dyana/cli.py b/dyana/cli.py
index 821584a..6717cf5 100644
--- a/dyana/cli.py
+++ b/dyana/cli.py
@@ -98,8 +98,8 @@ def trace(
 ) -> None:
     try:
         # disable GPU on non-Linux systems
-        if not no_gpu and platform_pkg.system() != "Linux":
-            no_gpu = True
+        # if not no_gpu and platform_pkg.system() != "Linux":
+        #    no_gpu = True
 
         # check if policy is either a file or a directory
         if policy and not policy.exists():
diff --git a/dyana/loaders/automodel/settings.yml b/dyana/loaders/automodel/settings.yml
index 8f6a05d..4cde030 100644
--- a/dyana/loaders/automodel/settings.yml
+++ b/dyana/loaders/automodel/settings.yml
@@ -1,5 +1,6 @@
 description: Loads and profiles machine learning models compatible with AutoModel and AutoTokenizer.
 
+gpu: true
 build_args:
   extra-requirements: EXTRA_REQUIREMENTS
 
diff --git a/dyana/loaders/loader.py b/dyana/loaders/loader.py
index 15abded..c32871c 100644
--- a/dyana/loaders/loader.py
+++ b/dyana/loaders/loader.py
@@ -220,7 +220,7 @@ def run(self, allow_network: bool = False, allow_gpus: bool = True, allow_volume
                 volumes,
                 environment=environment,
                 allow_network=allow_network,
-                allow_gpus=allow_gpus,
+                allow_gpus=allow_gpus and self.settings.gpu,
                 allow_volume_write=allow_volume_write,
             )
             self.container_id = self.container.id
diff --git a/dyana/loaders/lora/settings.yml b/dyana/loaders/lora/settings.yml
index 6e7bf99..ba55335 100644
--- a/dyana/loaders/lora/settings.yml
+++ b/dyana/loaders/lora/settings.yml
@@ -1,5 +1,7 @@
 description: Loads LoRA adapters via PEFT.
 
+gpu: true
+
 build_args:
   extra-requirements: EXTRA_REQUIREMENTS
 
diff --git a/dyana/loaders/ollama/settings.yml b/dyana/loaders/ollama/settings.yml
index ae6e88c..1c2b934 100644
--- a/dyana/loaders/ollama/settings.yml
+++ b/dyana/loaders/ollama/settings.yml
@@ -1,5 +1,7 @@
 description: Loads and profiles models via an Ollama server. Local models on the host machine are mounted and shared with the container.
 
+gpu: true
+
 volumes:
   # on macOS
   - host: ~/.ollama/models
diff --git a/dyana/loaders/settings.py b/dyana/loaders/settings.py
index 41c22cc..305ea79 100644
--- a/dyana/loaders/settings.py
+++ b/dyana/loaders/settings.py
@@ -32,6 +32,7 @@ class LoaderSettings(BaseModel):
     build_args: dict[str, str] | None = None
     args: list[LoaderArgument] | None = None
     network: bool | None = False
+    gpu: bool = False
     volumes: list[Volume] | None = None
     examples: list[Example] | None = None