From 490b70d89a6110d578443d240fe65ed6137da0fb Mon Sep 17 00:00:00 2001 From: Allen Chen Date: Thu, 10 Sep 2026 11:17:13 -0700 Subject: [PATCH] Require token authentication for Unsloth Jupyter --- README.md | 12 +++++++- docker-compose.yml | 3 +- scripts/configure_jupyter.py | 33 ++++++++++++++++++++ scripts/entrypoint-passwordless.sh | 40 ++++--------------------- scripts/test_jupyter_auth.py | 48 ++++++++++++++++++++++++++++++ 5 files changed, 99 insertions(+), 37 deletions(-) create mode 100644 scripts/configure_jupyter.py create mode 100644 scripts/test_jupyter_auth.py diff --git a/README.md b/README.md index 77af493..1e196bd 100644 --- a/README.md +++ b/README.md @@ -154,10 +154,20 @@ docker compose up -d #### Core Settings - `CONTAINER_NAME`: Container name (default: `unsloth-notebook`) - `JUPYTER_PORT`: Jupyter port (default: `8888`) -- `JUPYTER_PASSWORD`: Jupyter password protection +- `JUPYTER_TOKEN`: Optional login token; a secure random token is generated on each start when unset or empty. - `SSH_HOST_PORT`: SSH access port (default: `2222`) #### Authentication +Jupyter requires a token. Retrieve it locally with: + +```bash +docker compose exec unsloth-jupyter cat /home/unsloth/.jupyter/jupyter_token +``` + +Paste the token into the Jupyter login page. The token and configuration files are +readable only by their owner. An empty `JUPYTER_TOKEN` does not disable authentication. +The health check uses the public login page and does not include the token. + - `HF_TOKEN`: Hugging Face token for private models - `WANDB_API_KEY`: Weights & Biases API key - `SSH_KEY`: SSH public key for container access diff --git a/docker-compose.yml b/docker-compose.yml index 00c6395..84c9613 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -31,7 +31,6 @@ services: environment: # Unsloth / Jupyter JUPYTER_PORT: ${JUPYTER_PORT:-8888} - JUPYTER_PASSWORD: ${JUPYTER_PASSWORD:-} JUPYTER_TOKEN: ${JUPYTER_TOKEN:-} SSH_KEY: ${SSH_KEY:-} USER_PASSWORD: ${USER_PASSWORD:-unsloth2024} @@ -98,7 +97,7 @@ services: command: ${STARTUP_COMMAND:-/workspace/scripts/start-jupyter.sh} healthcheck: - test: ["CMD", "curl", "-f", "http://localhost:${JUPYTER_PORT:-8888}/api"] + test: ["CMD", "curl", "-f", "http://localhost:${JUPYTER_PORT:-8888}/login"] interval: 30s timeout: 10s retries: 3 diff --git a/scripts/configure_jupyter.py b/scripts/configure_jupyter.py new file mode 100644 index 0000000..ae9a954 --- /dev/null +++ b/scripts/configure_jupyter.py @@ -0,0 +1,33 @@ +"""Write authenticated Jupyter settings without printing credentials.""" + +import os +from pathlib import Path +import secrets + + +def configure_jupyter(config_dir: Path) -> None: + token = os.environ.get("JUPYTER_TOKEN") or secrets.token_urlsafe(32) + port = int(os.environ.get("JUPYTER_PORT", "8888")) + config_dir.mkdir(parents=True, exist_ok=True) + config = f''' +c.ServerApp.allow_root = False +c.ServerApp.allow_remote_access = True +c.ServerApp.open_browser = False +c.ServerApp.ip = "0.0.0.0" +c.ServerApp.port = {port} +c.ServerApp.notebook_dir = "/workspace" +c.ServerApp.terminado_settings = {{"shell_command": ["/bin/bash", "-l"]}} +c.IdentityProvider.token = {token!r} +''' + for name, content in (("jupyter_lab_config.py", config), ("jupyter_token", token)): + path = config_dir / name + # Protect both new files and files left behind by earlier starts. + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as stream: + os.fchmod(stream.fileno(), 0o600) + stream.write(content) + + +if __name__ == "__main__": + configure_jupyter(Path("/home/unsloth/.jupyter")) + print("Jupyter authentication enabled. Retrieve the token from ~/.jupyter/jupyter_token.") diff --git a/scripts/entrypoint-passwordless.sh b/scripts/entrypoint-passwordless.sh index 1c57cc0..e5fc581 100644 --- a/scripts/entrypoint-passwordless.sh +++ b/scripts/entrypoint-passwordless.sh @@ -1,8 +1,8 @@ #!/bin/bash set -e -# Custom Unsloth Entrypoint - Passwordless Jupyter -# This replaces the default Unsloth entrypoint to disable Jupyter authentication +# Custom Unsloth Entrypoint - authenticated Jupyter +# The historical filename is retained for existing Compose configurations. # Mount NVIDIA driver libraries (works across different driver versions) # bash /workspace/scripts/mount-nvidia-libs.sh @@ -12,7 +12,7 @@ set -e # export CUDA_HOME="/usr/local/cuda-12.8" echo "Exporting environment variables for SSH sessions..." -printenv | grep -E '^HF_|^CUDA_|^NCCL_|^JUPYTER_|^SSH_|^PUBLIC_|^USER_|^UNSLOTH_|^PATH=|^LD_LIBRARY_PATH=' | \ +printenv | grep -E '^HF_|^CUDA_|^NCCL_|^JUPYTER_PORT=|^SSH_|^PUBLIC_|^USER_|^UNSLOTH_|^PATH=|^LD_LIBRARY_PATH=' | \ sed 's/^\([^=]*\)=\(.*\)$/export \1="\2"/' > /tmp/unsloth_environment # Also explicitly set LD_LIBRARY_PATH and CUDA_HOME in bashrc @@ -84,37 +84,9 @@ else echo "✓ SSH host keys already exist and appear valid" fi -# Configure Jupyter - PASSWORDLESS MODE -echo "Generating Jupyter configuration (PASSWORDLESS MODE)..." -mkdir -p /home/unsloth/.jupyter - -python3 << 'EOFPYTHON' -import os - -config_file = '/home/unsloth/.jupyter/jupyter_lab_config.py' - -config_content = f''' -c.ServerApp.allow_root = False -c.ServerApp.allow_remote_access = True -c.ServerApp.open_browser = False -c.ServerApp.ip = "0.0.0.0" -c.ServerApp.port = {os.getenv('JUPYTER_PORT', '8888')} -c.ServerApp.notebook_dir = "/workspace" -c.ServerApp.terminado_settings = {{"shell_command": ["/bin/bash", "-l"]}} -c.ServerApp.allow_origin = "*" - -# Disable authentication - passwordless access -c.ServerApp.token = "" -c.ServerApp.password = "" -c.IdentityProvider.token = "" -''' - -with open(config_file, 'w') as f: - f.write(config_content) - -print(f'✓ Jupyter configured for PASSWORDLESS ACCESS') -print(f'✓ Config written to {config_file}') -EOFPYTHON +# Generate authenticated Jupyter configuration before starting supervisord. +export JUPYTER_CONFIG_DIR=/home/unsloth/.jupyter +python3 /workspace/scripts/configure_jupyter.py # Create SSH run directory - sshd looks for /run/sshd specifically echo "Setting up SSH prerequisites..." diff --git a/scripts/test_jupyter_auth.py b/scripts/test_jupyter_auth.py new file mode 100644 index 0000000..5d935da --- /dev/null +++ b/scripts/test_jupyter_auth.py @@ -0,0 +1,48 @@ +import os +from pathlib import Path +from types import SimpleNamespace +import tempfile +import unittest +from unittest.mock import patch + +from configure_jupyter import configure_jupyter + + +class JupyterAuthTest(unittest.TestCase): + def read_config(self, directory): + config = SimpleNamespace(ServerApp=SimpleNamespace(), IdentityProvider=SimpleNamespace()) + exec((directory / "jupyter_lab_config.py").read_text(), {"c": config}) + return config + + def test_empty_or_missing_token_generates_private_credential(self): + for env in ({}, {"JUPYTER_TOKEN": ""}): + with self.subTest(env=env), tempfile.TemporaryDirectory() as tmp: + directory = Path(tmp) + with patch.dict(os.environ, env, clear=True): + configure_jupyter(directory) + first = self.read_config(directory).IdentityProvider.token + configure_jupyter(directory) + config = self.read_config(directory) + self.assertGreaterEqual(len(first), 43) + self.assertNotEqual(first, config.IdentityProvider.token) + self.assertFalse(hasattr(config.ServerApp, "allow_origin")) + self.assertEqual(config.IdentityProvider.token, (directory / "jupyter_token").read_text()) + for name in ("jupyter_token", "jupyter_lab_config.py"): + self.assertEqual((directory / name).stat().st_mode & 0o777, 0o600) + + def test_config_quotes_explicit_token_and_restricts_existing_files(self): + token = "quotes'\"\n; c.IdentityProvider.token = '' #" + with tempfile.TemporaryDirectory() as tmp: + directory = Path(tmp) + for name in ("jupyter_token", "jupyter_lab_config.py"): + (directory / name).touch(mode=0o644) + with patch.dict(os.environ, {"JUPYTER_TOKEN": token, "JUPYTER_PORT": "9999"}, clear=True): + configure_jupyter(directory) + config = self.read_config(directory) + self.assertEqual(config.IdentityProvider.token, token) + self.assertEqual(config.ServerApp.port, 9999) + self.assertEqual((directory / "jupyter_lab_config.py").stat().st_mode & 0o777, 0o600) + + +if __name__ == "__main__": + unittest.main()