baschno/buun-stack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(jupyterhub): make SecretStore singleton

Browse Source
This commit is contained in:
Masaki Yatsu
2025-08-31 22:33:22 +09:00
parent 106714d0ac
commit 972805aa65
2 changed files with 37 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
jupyterhub/justfile
View File

@@ -3,9 +3,9 @@ set fallback := true
export JUPYTERHUB_NAMESPACE := env("JUPYTERHUB_NAMESPACE", "jupyter")
export JUPYTERHUB_CHART_VERSION := env("JUPYTERHUB_CHART_VERSION", "4.2.0")
export JUPYTERHUB_OIDC_CLIENT_ID := env("JUPYTERHUB_OIDC_CLIENT_ID", "jupyterhub")
export JUPYTERHUB_ENABLE_NFS_PV := env("JUPYTERHUB_ENABLE_NFS_PV", "")
export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "false")
export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-5")
export JUPYTERHUB_NFS_PV_ENABLED := env("JUPYTERHUB_NFS_PV_ENABLED", "")
export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "")
export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-6")
export KERNEL_IMAGE_BUUN_STACK_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_REPOSITORY", "buun-stack-notebook")
export KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY", "buun-stack-cuda-notebook")
export JUPYTER_PROFILE_MINIMAL_ENABLED := env("JUPYTER_PROFILE_MINIMAL_ENABLED", "false")
@@ -64,14 +64,14 @@ install:
export JUPYTER_PYTHON_KERNEL_TAG=${JUPYTER_PYTHON_KERNEL_TAG}
export JUPYTER_FSGID=${JUPYTER_FSGID:-100}
export PVC_NAME=""
if [ -z "${JUPYTERHUB_ENABLE_NFS_PV}" ]; then
if [ -z "${JUPYTERHUB_NFS_PV_ENABLED}" ]; then
if gum confirm "Are you going to use NFS PV?"; then
JUPYTERHUB_ENABLE_NFS_PV=true
JUPYTERHUB_NFS_PV_ENABLED=true
else
JUPYTERHUB_ENABLE_NFS_PV=false
JUPYTERHUB_NFS_PV_ENABLED=false
fi
fi
if [ "${JUPYTERHUB_ENABLE_NFS_PV}" = "true" ]; then
if [ "${JUPYTERHUB_NFS_PV_ENABLED}" = "true" ]; then
if ! helm status longhorn -n ${LONGHORN_NAMESPACE} &>/dev/null; then
echo "Longhorn is not installed. Please install Longhorn first." >&2
exit 1
@@ -104,7 +104,13 @@ install:
# wait deployments manually because `helm upgrade --wait` does not work for JupyterHub
just k8s::wait-deployments-ready ${JUPYTERHUB_NAMESPACE} hub proxy
# Setup Vault integration if enabled
if [ -z "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" ]; then
if gum confirm "Are you going to enable Vault integration?"; then
JUPYTERHUB_VAULT_INTEGRATION_ENABLED=true
else
JUPYTERHUB_VAULT_INTEGRATION_ENABLED=false
fi
fi
if [ "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" = "true" ]; then
just setup-vault-jwt-auth
fi

33
python-package/buunstack/secrets.py
View File

@@ -30,6 +30,9 @@ class SecretStore:
automatic OIDC token refresh via Keycloak integration and provides both
manual and background token management options.
This class implements the singleton pattern to ensure only one instance
exists per user session, preventing duplicate background refresh threads.
Attributes
----------
auto_token_refresh : bool
@@ -58,6 +61,14 @@ class SecretStore:
'sk-123'
"""
_instance = None
_initialized = False
def __new__(cls, *args, **kwargs):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(
self,
auto_token_refresh: bool = True,
@@ -67,18 +78,22 @@ class SecretStore:
"""
Initialize SecretStore with authentication and configuration.
Note: Due to singleton pattern, parameters are only used on the first
instantiation. Subsequent calls return the existing instance with
its original configuration.
Parameters
----------
auto_token_refresh : bool, optional
Enable automatic token refresh using Keycloak OIDC, by default True.
Requires KEYCLOAK_HOST, KEYCLOAK_REALM, and JUPYTERHUB_OIDC_REFRESH_TOKEN
environment variables.
environment variables. Only used on first instantiation.
refresh_buffer_seconds : int, optional
Seconds before token expiry to trigger refresh, by default 300.
Only used when auto_token_refresh is True.
Only used when auto_token_refresh is True. Only used on first instantiation.
background_refresh_interval : int, optional
Seconds between background refresh checks, by default 1800.
Only used when background refresh is started.
Only used when background refresh is started. Only used on first instantiation.
Raises
------
@@ -105,37 +120,33 @@ class SecretStore:
... background_refresh_interval=3600
... )
"""
if self._initialized:
return
self.auto_token_refresh = auto_token_refresh
self.refresh_buffer_seconds = refresh_buffer_seconds
self.background_refresh_interval = background_refresh_interval
# User and environment info
self.username = os.getenv("JUPYTERHUB_USER")
self.vault_addr = os.getenv("VAULT_ADDR")
# Keycloak configuration (only needed if auto_token_refresh is enabled)
if self.auto_token_refresh:
self.keycloak_host = os.getenv("KEYCLOAK_HOST")
self.keycloak_realm = os.getenv("KEYCLOAK_REALM")
self.keycloak_client_id = os.getenv("KEYCLOAK_CLIENT_ID", "jupyterhub")
self.refresh_token = os.getenv("JUPYTERHUB_OIDC_REFRESH_TOKEN")
# Token management
self.access_token = os.getenv("JUPYTERHUB_OIDC_ACCESS_TOKEN")
self.token_expiry = (
self._get_token_expiry(self.access_token) if self.access_token else None
)
# Initialize Vault client
self.client = hvac.Client(url=self.vault_addr, verify=False)
# Background refresher
self._background_refresher = None
# Authenticate initially
self._authenticate_vault()
# Set base path for user storage
self.base_path = f"jupyter/users/{self.username}"
logger.info(f"SecretStore initialized for user: {self.username}")
@@ -146,6 +157,8 @@ class SecretStore:
if self.auto_token_refresh and self.token_expiry:
logger.info(f"Token expires at: {self.token_expiry}")
self._initialized = True
def _get_token_expiry(self, token: str) -> datetime | None:
"""Extract expiry time from JWT token"""
if not token: