Add background provisioning with JupyterHub integration

yacchin1205 · yacchin1205 · commit 0bb2902c132d · 2025-10-22T07:25:24.000+09:00
diff --git a/repo2docker/contentproviders/rdm/README.md b/repo2docker/contentproviders/rdm/README.md
@@ -55,3 +55,143 @@ paths:
       source: $default_storage_path
       target: .
 ```
+
+## Running provision.sh with JupyterHub
+
+When deploying repo2docker-built images with JupyterHub, you can automatically execute the `provision.sh` script at container startup to provision RDM data.
+
+### Background Execution to Avoid Timeout
+
+Since copying large datasets may take time and cause JupyterHub spawn timeout, the `provision.sh` script supports background execution mode. When called with command-line arguments, it will:
+
+1. Start provisioning (copy/link operations) in the background
+2. Immediately execute the passed command (e.g., `jupyterhub-singleuser`)
+
+This allows the JupyterHub server to start while data provisioning continues in the background.
+
+### JupyterHub Configuration
+
+Configure your JupyterHub spawner to execute `provision.sh` if it exists:
+
+#### KubeSpawner Example
+
+```python
+# In jupyterhub_config.py
+c.KubeSpawner.cmd = [
+    'bash', '-c',
+    '''
+    set -e
+
+    # Find and execute provision.sh if it exists
+    for path in \
+        "${REPO_DIR}/binder/provision.sh" \
+        "${REPO_DIR}/.binder/provision.sh" \
+        "$HOME/binder/provision.sh" \
+        "$HOME/.binder/provision.sh"; do
+
+        if [ -f "$path" ]; then
+            echo "[provision-wrapper] Executing: $path" >&2
+            exec bash "$path" "$@"
+        fi
+    done
+
+    # No provision.sh found, start normally
+    exec "$@"
+    ''',
+    '--', 'jupyterhub-singleuser'
+]
+```
+
+#### DockerSpawner Example
+
+```python
+# In jupyterhub_config.py
+c.DockerSpawner.cmd = [
+    'bash', '-c',
+    '''
+    set -e
+    for path in \
+        "${REPO_DIR}/binder/provision.sh" \
+        "${REPO_DIR}/.binder/provision.sh" \
+        "$HOME/binder/provision.sh" \
+        "$HOME/.binder/provision.sh"; do
+        [ -f "$path" ] && exec bash "$path" "$@"
+    done
+    exec "$@"
+    ''',
+    '--', 'jupyterhub-singleuser'
+]
+```
+
+### How provision.sh Works
+
+The generated `provision.sh` script accepts command-line arguments and has the following structure:
+
+```bash
+#!/bin/bash
+set -e
+
+# Run provisioning in background
+{
+    # Copy and link operations
+    mkdir -p './target/path/'
+    cp -fr '/mnt/rdm/storage/data/'* './target/path/'
+    ln -s '/mnt/rdm/large-data/' './data'
+} &
+
+# Execute passed command if provided
+if [ $# -gt 0 ]; then
+    exec "$@"
+fi
+```
+
+### Volume Mounts
+
+Ensure that RDM storage is mounted at `/mnt/rdm/` in the container. Configure your spawner accordingly:
+
+```python
+# KubeSpawner example
+c.KubeSpawner.volumes = [
+    {
+        'name': 'rdm-storage',
+        'persistentVolumeClaim': {
+            'claimName': 'rdm-pvc'
+        }
+    }
+]
+
+c.KubeSpawner.volume_mounts = [
+    {
+        'name': 'rdm-storage',
+        'mountPath': '/mnt/rdm/'
+    }
+]
+```
+
+**Note**: If `/mnt/rdm/` does not exist but `/mnt/rdms/{project_id}/` is available, `provision.sh` will automatically create a symlink from `/mnt/rdm` to `/mnt/rdms/{project_id}` at startup.
+
+### Monitoring Provisioning Progress
+
+Users can check the provisioning progress from within the Jupyter environment:
+
+```bash
+# View the provisioning log in real-time
+tail -f /tmp/provision.log
+
+# Check if provisioning is complete
+grep "completed" /tmp/provision.log
+```
+
+The log file `/tmp/provision.log` contains:
+- Start and completion timestamps
+- Each copy/link operation with source and target paths
+- Detailed command output (from `set -x`)
+- Any errors that occur during provisioning
+
+### Notes
+
+- The `REPO_DIR` environment variable points to the repository directory (default: `/home/jovyan`)
+- Provisioning runs in the background, so large data copies won't block JupyterHub startup
+- Symbolic links are created immediately and are available right away
+- Check `/tmp/provision.log` for provisioning progress and errors
+- Container logs will show when provisioning starts and how to monitor it
diff --git a/repo2docker/contentproviders/rdm/__init__.py b/repo2docker/contentproviders/rdm/__init__.py
@@ -255,6 +255,7 @@ async def _fetch_binder(
         for binder_output_dir in binder_output_dirs:
             provisioner.save_provision_script(
                 os.path.join(binder_output_dir, "provision.sh"),
+                project.id,
                 mnt_rdm_dir,
             )
 
diff --git a/repo2docker/contentproviders/rdm/provisioner.py b/repo2docker/contentproviders/rdm/provisioner.py
@@ -57,50 +57,76 @@ async def add_link_mapping(self, path_mapping: PathMapping):
         source = await self._resolve_source(path_mapping)
         self._link_mappings.append((path_mapping, source))
 
-    def save_provision_script(self, script_path: str, source_mount_dir='/mnt/rdm/'):
+    def save_provision_script(self, script_path: str, project_id: str, source_mount_dir='/mnt/rdm/'):
         """Save the provision script to the specified path."""
         with open(script_path, "w") as f:
             f.write("#!/bin/bash\n")
-            f.write("set -xe\n")
+            f.write("set -e\n\n")
+            f.write("# Ensure /mnt/rdm exists or create symlink to project-specific directory\n")
+            f.write("if [ ! -e /mnt/rdm ]; then\n")
+            f.write(f"    PROJECT_DIR=/mnt/rdms/{shlex.quote(project_id)}\n")
+            f.write("    for i in 1 2 4; do\n")
+            f.write("        if [ -d \"$PROJECT_DIR\" ]; then\n")
+            f.write("            ln -s \"$PROJECT_DIR\" /mnt/rdm\n")
+            f.write("            break\n")
+            f.write("        fi\n")
+            f.write("        echo \"Waiting for $PROJECT_DIR to be available... (retry in ${i}s)\" >&2\n")
+            f.write("        sleep $i\n")
+            f.write("    done\n")
+            f.write("fi\n\n")
+            f.write("PROVISION_LOG=\"/tmp/provision.log\"\n\n")
+            f.write("# Run provisioning in background\n")
+            f.write("{\n")
+            f.write("    echo \"[provision] Starting RDM data provisioning at $(date)...\"\n")
+            f.write("    set -x\n")
             for path_mapping, source in self._copy_mappings:
                 source_path = os.path.join(
                     source_mount_dir,
                     path_mapping.get_source(self._default_storage_path)
                 )
                 target_path = path_mapping.get_target()
+                f.write(f"    echo \"[provision] Copying {shlex.quote(source_path)} to {shlex.quote(target_path)}...\"\n")
                 if target_path != "./" and target_path.endswith("/"):
                     # target is directory
                     if target_path.strip("/") != "." and target_path.strip("/") != "":
-                        f.write(f"mkdir -p {shlex.quote(target_path)}\n")
+                        f.write(f"    mkdir -p {shlex.quote(target_path)}\n")
                 elif target_path != "./" and "/" in target_path:
                     # target is subdir
                     parent_dir = os.path.dirname(target_path)
                     if parent_dir.strip("/") != "." and parent_dir.strip("/") != "":
-                        f.write(f"mkdir -p {shlex.quote(parent_dir)}\n")
+                        f.write(f"    mkdir -p {shlex.quote(parent_dir)}\n")
                 if source.path.endswith("/"):
                     # folder
                     if target_path != "." and not target_path.endswith("/"):
                         target_path += "/"
                     if not source_path.endswith("/"):
                         source_path += "/"
                     if target_path.strip("/") != "." and target_path.strip("/") != "":
-                        f.write(f"mkdir -p {shlex.quote(target_path)}\n")
-                    f.write(f"cp -fr {shlex.quote(source_path)}* {shlex.quote(target_path)}\n")
+                        f.write(f"    mkdir -p {shlex.quote(target_path)}\n")
+                    f.write(f"    cp -fr {shlex.quote(source_path)}* {shlex.quote(target_path)}\n")
                 else:
                     # file
-                    f.write(f"cp {shlex.quote(source_path)} {shlex.quote(target_path)}\n")
+                    f.write(f"    cp {shlex.quote(source_path)} {shlex.quote(target_path)}\n")
             for path_mapping, source in self._link_mappings:
                 source_path = os.path.join(
                     source_mount_dir,
                     path_mapping.get_source(self._default_storage_path)
                 )
                 target_path = path_mapping.get_target()
+                f.write(f"    echo \"[provision] Linking {shlex.quote(source_path)} to {shlex.quote(target_path)}...\"\n")
                 if target_path != "./" and "/" in target_path.strip("/"):
                     # target is subdir
                     parent_dir = os.path.dirname(target_path)
                     if parent_dir.strip("/") != "." and parent_dir.strip("/") != "":
-                        f.write(f"mkdir -p {shlex.quote(parent_dir)}\n")
-                f.write(f"ln -s {shlex.quote(source_path)} {shlex.quote(target_path)}\n")
+                        f.write(f"    mkdir -p {shlex.quote(parent_dir)}\n")
+                f.write(f"    ln -s {shlex.quote(source_path)} {shlex.quote(target_path)}\n")
+            f.write("    echo \"[provision] RDM data provisioning completed at $(date).\"\n")
+            f.write("} > \"${PROVISION_LOG}\" 2>&1 &\n\n")
+            f.write("echo \"[provision] Provisioning started in background. Check progress: tail -f ${PROVISION_LOG}\" >&2\n\n")
+            f.write("# Execute passed command if provided\n")
+            f.write("if [ $# -gt 0 ]; then\n")
+            f.write("    exec \"$@\"\n")
+            f.write("fi\n")
 
     async def _resolve_source(self, path_mapping: PathMapping):
         """Validate the path mapping."""
diff --git a/tests/unit/contentproviders/test_rdm.py b/tests/unit/contentproviders/test_rdm.py
@@ -416,6 +416,7 @@ async def mock_storage(name):
 
             fake_project_obj = MagicMock(storages=AsyncIterator([fake_storage]))
             fake_project_obj.storage = mock_storage
+            fake_project_obj.id = "x1234"
             fake_project.return_value = fake_project_obj
 
             binder_dir = os.path.join(d, "binder")
@@ -469,9 +470,20 @@ async def mock_resolve_source(self, path_mapping):
 
                 # Verify script content
                 assert "#!/bin/bash" in script_content
-                assert "set -xe" in script_content
+                assert "set -e" in script_content
+                # Verify /mnt/rdm symlink creation with retry logic for project-specific directory
+                assert "if [ ! -e /mnt/rdm ]; then" in script_content
+                assert "PROJECT_DIR=/mnt/rdms/x1234" in script_content
+                assert "for i in 1 2 4; do" in script_content
+                assert 'ln -s "$PROJECT_DIR" /mnt/rdm' in script_content
+                assert "Waiting for $PROJECT_DIR to be available" in script_content
+                assert 'PROVISION_LOG="/tmp/provision.log"' in script_content
                 assert "cp -fr /mnt/rdm/osfstorage/data/* ./dataset/" in script_content
                 assert "ln -s /mnt/rdm/external_storage/large_files ./external" in script_content
+                # Verify background execution and command passing
+                assert "} > \"${PROVISION_LOG}\" 2>&1 &" in script_content
+                assert 'if [ $# -gt 0 ]; then' in script_content
+                assert 'exec "$@"' in script_content
 
 
 def test_fetch_with_binder_but_no_paths_yaml():
@@ -501,6 +513,7 @@ async def mock_storage(name):
 
             fake_project_obj = MagicMock(storages=AsyncIterator([fake_storage]))
             fake_project_obj.storage = mock_storage
+            fake_project_obj.id = "x1234"
             fake_project.return_value = fake_project_obj
 
             # Mock Provisioner._resolve_source to avoid storage validation
@@ -529,10 +542,14 @@ async def mock_resolve_source(self, path_mapping):
 
                 # Verify default mapping is added (copy entire storage to current directory)
                 assert "#!/bin/bash" in script_content
-                assert "set -xe" in script_content
+                assert "set -e" in script_content
+                # Verify /mnt/rdm symlink creation with retry logic
+                assert "PROJECT_DIR=/mnt/rdms/x1234" in script_content
+                assert 'ln -s "$PROJECT_DIR" /mnt/rdm' in script_content
                 assert "cp -fr /mnt/rdm/osfstorage/* ." in script_content
-                # Should not have any link commands
-                assert "ln -s" not in script_content
+                # Should not have any user-defined link commands in the background block
+                # (only the /mnt/rdm setup link should exist)
+                assert script_content.count("ln -s") == 1
 
 
 def test_fetch_with_empty_paths_and_override():
@@ -561,6 +578,7 @@ async def mock_storage(name):
 
             fake_project_obj = MagicMock(storages=AsyncIterator([fake_storage]))
             fake_project_obj.storage = mock_storage
+            fake_project_obj.id = "x1234"
             fake_project.return_value = fake_project_obj
 
             binder_dir = os.path.join(d, "binder")
@@ -593,17 +611,21 @@ async def mock_resolve_source(self, path_mapping):
                 with open(provision_script_path, 'r') as f:
                     script_content = f.read()
 
-                # Verify only shebang and set -xe are present, no copy or link commands
+                # Verify script structure without user-defined copy or link commands
                 assert "#!/bin/bash" in script_content
-                assert "set -xe" in script_content
+                assert "set -e" in script_content
+                # Verify /mnt/rdm symlink creation with retry logic
+                assert "PROJECT_DIR=/mnt/rdms/x1234" in script_content
+                assert 'ln -s "$PROJECT_DIR" /mnt/rdm' in script_content
                 # Should not have any copy commands
                 assert "cp -fr" not in script_content
                 assert "cp " not in script_content
-                # Should not have any link commands
-                assert "ln -s" not in script_content
-                # The script should only have 2 lines
-                lines = [line for line in script_content.strip().split('\n') if line]
-                assert len(lines) == 2, f"Expected 2 lines, got {len(lines)}: {lines}"
+                # Should not have any user-defined link commands
+                # (only the /mnt/rdm setup link should exist)
+                assert script_content.count("ln -s") == 1
+                # But should have background execution structure
+                assert 'PROVISION_LOG="/tmp/provision.log"' in script_content
+                assert "} > \"${PROVISION_LOG}\" 2>&1 &" in script_content
 
 
 def test_rdmurl_project_id():

Original file line number	Diff line number	Diff line change
`@@ -255,6 +255,7 @@ async def _fetch_binder(`
`255`	`255`	`for binder_output_dir in binder_output_dirs:`
`256`	`256`	`provisioner.save_provision_script(`
`257`	`257`	`os.path.join(binder_output_dir, "provision.sh"),`
	`258`	`+ project.id,`
`258`	`259`	`mnt_rdm_dir,`
`259`	`260`	`)`
`260`	`261`