Run nvidia_runtime_is_available check using agents image (#93)

* run same image when checking for nvidia runtime availability
supervisely · Dec 4, 2024 · 5775770 · 5775770
1 parent 1646dba
commit 5775770
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/agent/worker/agent_utils.py b/agent/worker/agent_utils.py
@@ -1069,16 +1069,19 @@ def restart_agent(
 
 def nvidia_runtime_is_available():
     docker_api = docker.from_env()
-    image = constants.DEFAULT_APP_DOCKER_IMAGE()
+    container_info = get_container_info()
+    image = container_info.get("Image")
     try:
         docker_api.containers.run(
             image,
+            entrypoint="",
             command="nvidia-smi",
             runtime="nvidia",
             remove=True,
         )
         return True
     except Exception as e:
+        sly.logger.warning(f"Failed to run nvidia-smi:", exc_info=True)
         return False
 
 

diff --git a/agent/worker/constants.py b/agent/worker/constants.py
@@ -145,7 +145,7 @@ def TOKEN():
     _SUPERVISELY_AGENT_FILES: f"/opt/supervisely/agents/agent-files-{TOKEN()[:8]}",
     _SUPERVISELY_AGENT_FILES_CONTAINER: "/app/sly-files",
     _OFFLINE_MODE: False,
-    _DEFAULT_APP_DOCKER_IMAGE: "supervisely/base-py-sdk",
+    _DEFAULT_APP_DOCKER_IMAGE: "supervisely/base-py-sdk:latest",
     _AGENT_FILES_IN_APP_CONTAINER: "/agent-storage",
     _AUTO_CLEAN_INT_RANGE_DAYS: 7,
     _REQUESTS_CA_BUNDLE_DIR_CONTAINER: "/sly_certs",