Require token authentication for Unsloth Jupyter

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

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

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.")

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..."

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()