Gpu device info (#29)

* not tested * gpu_info moved to hardware_info * agent_info removed from agent_info --------- Co-authored-by: Maxim Kolomeychenko <[email protected]> Co-authored-by: Umar Sadekov <[email protected]>
supervisely · Jun 15, 2023 · 7bc13e6 · 7bc13e6
1 parent b500446
commit 7bc13e6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -58,6 +58,8 @@ RUN pip install httpx
 RUN pip install supervisely==6.72.7
 RUN pip install requests-toolbelt>=1.0.0
 
+RUN pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html
+
 COPY . /workdir
 
 #ENV PYTHONPATH /workdir:/workdir/src:/workdir/supervisely_lib/worker_proto:$PYTHONPATH

diff --git a/agent/worker/agent.py b/agent/worker/agent.py
@@ -9,6 +9,7 @@
 import os
 import supervisely_lib as sly
 import uuid
+import torch
 
 from worker import constants
 from worker.task_factory import create_task, is_task_type
@@ -115,6 +116,20 @@ def agent_connect_initially(self):
             config = {}
             self.logger.warn("Docker container info unavailable, agent is running in debug VENV")
 
+        gpu_info = None
+        try:
+            gpu_info = {}
+            gpu_info["is_available"] = torch.cuda.is_available()
+            if gpu_info["is_available"]:
+                gpu_info["device_count"] = torch.cuda.device_count()
+                gpu_info["device_names"] = []
+                for idx in range(gpu_info["device_count"]):
+                    gpu_info["device_names"].append(torch.cuda.get_device_name(idx))
+        except Exception as e:
+            self.logger.warning(repr(e))
+
+        hw_info["gpuinfo"] = gpu_info
+
         self.agent_info = {
             "agent_id_from_storage_path": constants.AGENT_ID(),
             "hardware_info": hw_info,
@@ -128,7 +143,7 @@ def agent_connect_initially(self):
                 constants._DOCKER_NET: constants.DOCKER_NET(),
             },
         }
-
+        # @todo: add GPU device info
         resp = self.api.simple_request(
             "AgentConnected",
             sly.api_proto.ServerInfo,

diff --git a/requirements.txt b/requirements.txt
@@ -12,4 +12,5 @@ docker==3.3.0
 black
 python-slugify==6.1.2
 requests==2.28.1
-urllib3==1.26.11
+urllib3==1.26.11
+torch==1.7.1