refactor: add logging for missing nvidia-smi command in get_cuda_vers…

…ion + ifelse logic return None
FocoosAI · Dec 24, 2024 · 87c82d4 · 87c82d4
1 parent cf70cda
commit 87c82d4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/focoos/utils/system.py b/focoos/utils/system.py
@@ -10,6 +10,9 @@
 
 from focoos.config import FOCOOS_CONFIG
 from focoos.ports import GPUInfo, SystemInfo
+from focoos.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class HttpClient:
@@ -155,11 +158,8 @@ def get_cuda_version() -> Optional[str]:
                     cuda_version = line.split(":")[-1].strip()
                     cuda_version = cuda_version.split()[0]
                     return cuda_version
-            return None
-        else:
-            return None
-    except FileNotFoundError:
-        return None
+    except FileNotFoundError as err:
+        logger.warning("nvidia-smi command not found: %s", err)
 
 
 def get_gpu_name() -> Optional[str]:

diff --git a/tests/test_system.py b/tests/test_system.py
@@ -1,7 +1,5 @@
 from unittest.mock import MagicMock, patch
 
-import pytest
-
 from focoos.ports import SystemInfo
 from focoos.utils.system import (
     HttpClient,