From 800c259bee2cae5b00a672e8b1b8fd9fbc8f6eea Mon Sep 17 00:00:00 2001
From: stellarshenson <konrad.jelen@gmail.com>
Date: Thu, 7 Aug 2025 20:08:36 +0200
Subject: [PATCH] updated with the nvidia autodetection

---
 compose.yml                               |  3 ++-
 config/jupyterhub_config.py               | 18 ++++++++++++++++--
 services/jupyterhub/Dockerfile.jupyterhub |  1 +
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/compose.yml b/compose.yml
index bb82c75..246afbb 100644
--- a/compose.yml
+++ b/compose.yml
@@ -53,12 +53,13 @@ services:
       - /var/run/docker.sock:/var/run/docker.sock:rw # docker socket to use spawner
       - jupyterhub_data:/data # database and cookie secrets
       - jupyterhub_certs:/mnt/certs # SSL certificates
+      - /var/run/docker.sock:/var/run/docker.sock:ro # for nvidia autodetection
     environment:
       - JUPYTERHUB_ADMIN=admin # this username will be a JupyterHub admin
       - DOCKER_NOTEBOOK_IMAGE=stellars/stellars-jupyterlab-ds:latest # jupyterlab image to spawn
       - DOCKER_NETWORK_NAME=jupyterhub-network # spawned containers will join this network
       - JUPYTERHUB_BASE_URL=/jupyterhub # default prefix
-      - ENABLE_GPU_SUPPORT=0 # this enables NVIDIA GPU support for images
+      - ENABLE_GPU_SUPPORT=2 # gpu status: 0 - disabled, 1 - enabled, 2 - auto-detect
       - ENABLE_JUPYTERHUB_SSL=0 # if using traefik - you do need direct SSL config
     labels:
       # Enable proxy support from Traefik
diff --git a/config/jupyterhub_config.py b/config/jupyterhub_config.py
index 37332f0..94ad1d5 100644
--- a/config/jupyterhub_config.py
+++ b/config/jupyterhub_config.py
@@ -6,17 +6,30 @@ import os
 import json
 import requests
 import nativeauthenticator
+import docker # for gpu autodetection
 
 c = get_config()  
 
+# NVIDIA GPU auto-detection
+NVIDIA_DETECTED = 0
+docker_client = docker.DockerClient(base_url='unix://var/run/docker.sock')
+docker_info = docker_client.info()
+if "nvidia" in docker_info.get("Runtimes", {}):
+    NVIDIA_DETECTED = 1
+
 # standard variables imported from env
 ENABLE_JUPYTERHUB_SSL =  int(os.environ.get("ENABLE_JUPYTERHUB_SSL", 1))
-ENABLE_GPU_SUPPORT= int(os.environ.get("ENABLE_GPU_SUPPORT", 0))
+ENABLE_GPU_SUPPORT= int(os.environ.get("ENABLE_GPU_SUPPORT", 2)) 
 DOCKER_NOTEBOOK_DIR = "/home/lab/workspace"
 JUPYTERHUB_BASE_URL = os.environ.get("JUPYTERHUB_BASE_URL")
 JUPYTERHUB_ADMIN = os.environ.get("JUPYTERHUB_ADMIN")
 NETWORK_NAME = os.environ["DOCKER_NETWORK_NAME"]
 
+# enable gpu autodetect and GPU found
+# gpu support: 0 - disabled, 1 - enabled, 2 - autodetect
+if ENABLE_GPU_SUPPORT == 2 and NVIDIA_DETECTED:
+    ENABLE_GPU_SUPPORT = 1 # means - gpu enabled
+
 # ensure that we are using SSL, it should be enabled by default
 if ENABLE_JUPYTERHUB_SSL == 1:
     c.JupyterHub.ssl_cert = '/mnt/certs/server.crt'
@@ -37,7 +50,8 @@ c.DockerSpawner.environment = {
      'ENABLE_SERVICE_GLANCES':1,
      'ENABLE_SERVICE_TENSORBOARD':1,
      'ENABLE_GPU_SUPPORT': ENABLE_GPU_SUPPORT,
-     'ENABLE_GPUSTAT': ENABLE_GPU_SUPPORT
+     'ENABLE_GPUSTAT': ENABLE_GPU_SUPPORT,
+     'NVIDIA_DETECTED': NVIDIA_DETECTED,
 }
 
 # configure access to GPU if possible
diff --git a/services/jupyterhub/Dockerfile.jupyterhub b/services/jupyterhub/Dockerfile.jupyterhub
index d9bc160..1b7e22d 100644
--- a/services/jupyterhub/Dockerfile.jupyterhub
+++ b/services/jupyterhub/Dockerfile.jupyterhub
@@ -37,6 +37,7 @@ COPY ./templates/certs /mnt/certs
 
 ## install dockerspawner, nativeauthenticator
 RUN pip install -U --no-cache-dir \
+    docker \
     dockerspawner \
     jupyterhub-nativeauthenticator