force use of "spawn" process creation method at the start of extension.initialize_settings

Author: andrii-i

jupyter_scheduler/extension.py

@@ -1,8 +1,6 @@
 import asyncio
 import multiprocessing
 
-multiprocessing.set_start_method("spawn", force=True)
-
 from jupyter_core.paths import jupyter_data_dir
 from jupyter_server.extension.application import ExtensionApp
 from jupyter_server.transutils import _i18n
@@ -72,6 +70,15 @@ def _db_url_default(self):
     )
 
     def initialize_settings(self):
+        # Forces new processes to not be forked on Linux.
+        # This is necessary because `asyncio.get_event_loop()` is bugged in
+        # forked processes in Python versions below 3.12. This method is
+        # called by `jupyter_core` by `nbconvert` in the default executor.
+
+        # See: https://github.com/python/cpython/issues/66285
+        # See also: https://github.com/jupyter/jupyter_core/pull/362
+        multiprocessing.set_start_method("spawn", force=True)
+
         super().initialize_settings()
 
         create_tables(self.db_url, self.drop_tables)

jupyter_scheduler/job_files_manager.py

@@ -1,4 +1,4 @@
-import multiprocessing as mp
+from multiprocessing import Process
 import os
 import random
 import tarfile
@@ -23,15 +23,7 @@ async def copy_from_staging(self, job_id: str, redownload: Optional[bool] = Fals
         output_filenames = self.scheduler.get_job_filenames(job)
         output_dir = self.scheduler.get_local_output_path(model=job, root_dir_relative=True)
 
-        # The MP context forces new processes to not be forked on Linux.
-        # This is necessary because `asyncio.get_event_loop()` is bugged in
-        # forked processes in Python versions below 3.12. This method is
-        # called by `jupyter_core` by `nbconvert` in the default executor.
-        #
-        # See: https://github.com/python/cpython/issues/66285
-        # See also: https://github.com/jupyter/jupyter_core/pull/362
-        mp_ctx = mp.get_context("spawn")
-        p = mp_ctx.Process(
+        p = Process(
             target=Downloader(
                 output_formats=job.output_formats,
                 output_filenames=output_filenames,

jupyter_scheduler/scheduler.py

@@ -1,4 +1,4 @@
-import multiprocessing as mp
+from multiprocessing import Process
 import os
 import random
 import shutil
@@ -481,15 +481,7 @@ def create_job(self, model: CreateJob) -> str:
             else:
                 self.copy_input_file(model.input_uri, staging_paths["input"])
 
-            # The MP context forces new processes to not be forked on Linux.
-            # This is necessary because `asyncio.get_event_loop()` is bugged in
-            # forked processes in Python versions below 3.12. This method is
-            # called by `jupyter_core` by `nbconvert` in the default executor.
-            #
-            # See: https://github.com/python/cpython/issues/66285
-            # See also: https://github.com/jupyter/jupyter_core/pull/362
-            mp_ctx = mp.get_context("spawn")
-            p = mp_ctx.Process(
+            p = Process(
                 target=self.execution_manager_class(
                     job_id=job.job_id,
                     staging_paths=staging_paths,

jupyter_scheduler/tests/test_job_files_manager.py

@@ -4,7 +4,7 @@
 import tarfile
 import time
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 
@@ -40,20 +40,13 @@ async def test_copy_from_staging():
         "input": "helloworld.ipynb",
     }
     output_dir = "jobs/1"
-
-    mock_scheduler = MagicMock()
-    mock_scheduler.get_job.return_value = job
-    mock_scheduler.get_staging_paths.return_value = staging_paths
-    mock_scheduler.get_local_output_path.return_value = output_dir
-    mock_scheduler.get_job_filenames.return_value = job_filenames
-
-    mock_context = MagicMock()
-    mock_process = MagicMock()
-    mock_context.Process.return_value = mock_process
-
     with patch("jupyter_scheduler.job_files_manager.Downloader") as mock_downloader:
-        with patch("jupyter_scheduler.scheduler.Scheduler", return_value=mock_scheduler):
-            with patch("multiprocessing.get_context", return_value=mock_context):
+        with patch("jupyter_scheduler.job_files_manager.Process") as mock_process:
+            with patch("jupyter_scheduler.scheduler.Scheduler") as mock_scheduler:
+                mock_scheduler.get_job.return_value = job
+                mock_scheduler.get_staging_paths.return_value = staging_paths
+                mock_scheduler.get_local_output_path.return_value = output_dir
+                mock_scheduler.get_job_filenames.return_value = job_filenames
                 manager = JobFilesManager(scheduler=mock_scheduler)
                 await manager.copy_from_staging(1)
 
@@ -65,7 +58,6 @@ async def test_copy_from_staging():
                     redownload=False,
                     include_staging_files=None,
                 )
-                mock_process.start.assert_called_once()
 
 
 HERE = Path(__file__).parent.resolve()

jupyter_scheduler/utils.py

@@ -1,4 +1,5 @@
 import json
+import multiprocessing
 import os
 import shutil
 from datetime import datetime, timezone
@@ -93,7 +94,7 @@ def copy_directory(
     exclude_files: Optional[List[str]] = [],
 ) -> List[str]:
     """Copies content of source_dir to destination_dir excluding exclude_files.
-    Returns a list of relative paths to copied files from destination_dir.
+    Returns a list of relative paths to copied files from destination_dir
     """
     copied_files = []
     for item in os.listdir(source_dir):