fix: replace multiprocessing.set_start_method('fork') with get_context('fork') (#478)

* fix: replace multiprocessing.set_start_method('fork') with get_context('fork')

- Use multiprocessing.get_context('fork') instead of set_start_method('fork')
  to avoid RuntimeError when other libraries set the start method first
- Route all multiprocessing calls through _fork_context in worker.py
- Default to gokart's WorkerSchedulerFactory in build() and run()

Closes #469

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* docs: update fork comment

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

gokart/build.py

@@ -211,6 +211,7 @@ def _build_task():
             task_lock_exception_raised.flag = False
             result = luigi.build(
                 [task],
+                worker_scheduler_factory=WorkerSchedulerFactory(),
                 local_scheduler=True,
                 detailed_summary=True,
                 log_level=logging.getLevelName(log_level),

gokart/run.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
+import logging
 import os
 import sys
 from logging import getLogger
 
 import luigi
 import luigi.cmdline
+import luigi.cmdline_parser
+import luigi.execution_summary
+import luigi.interface
 import luigi.retcodes
+import luigi.setup_logging
 from luigi.cmdline_parser import CmdlineParser
 
 import gokart
 import gokart.slack
+from gokart.build import WorkerSchedulerFactory
 from gokart.object_storage import ObjectStorage
 
 logger = getLogger(__name__)
@@ -80,6 +86,47 @@ def _try_to_send_event_summary_to_slack(slack_api: gokart.slack.SlackAPI | None,
     slack_api.send_snippet(comment=comment, title='event.txt', content=content)
 
 
+def _run_with_retcodes(argv):
+    """run_with_retcodes equivalent that uses gokart's WorkerSchedulerFactory."""
+    retcode_logger = logging.getLogger('luigi-interface')
+    with luigi.cmdline_parser.CmdlineParser.global_instance(argv):
+        retcodes = luigi.retcodes.retcode()
+
+    worker = None
+    try:
+        worker = luigi.interface._run(argv, worker_scheduler_factory=WorkerSchedulerFactory()).worker
+    except luigi.interface.PidLockAlreadyTakenExit:
+        sys.exit(retcodes.already_running)
+    except Exception:
+        env_params = luigi.interface.core()
+        luigi.setup_logging.InterfaceLogging.setup(env_params)
+        retcode_logger.exception('Uncaught exception in luigi')
+        sys.exit(retcodes.unhandled_exception)
+
+    with luigi.cmdline_parser.CmdlineParser.global_instance(argv):
+        task_sets = luigi.execution_summary._summary_dict(worker)
+        root_task = luigi.execution_summary._root_task(worker)
+        non_empty_categories = {k: v for k, v in task_sets.items() if v}.keys()
+
+    def has(status):
+        assert status in luigi.execution_summary._ORDERED_STATUSES
+        return status in non_empty_categories
+
+    codes_and_conds = (
+        (retcodes.missing_data, has('still_pending_ext')),
+        (retcodes.task_failed, has('failed')),
+        (retcodes.already_running, has('run_by_other_worker')),
+        (retcodes.scheduling_error, has('scheduling_error')),
+        (retcodes.not_run, has('not_run')),
+    )
+    expected_ret_code = max(code * (1 if cond else 0) for code, cond in codes_and_conds)
+
+    if expected_ret_code == 0 and root_task not in task_sets['completed'] and root_task not in task_sets['already_done']:
+        sys.exit(retcodes.not_run)
+    else:
+        sys.exit(expected_ret_code)
+
+
 def run(cmdline_args=None, set_retcode=True):
     cmdline_args = cmdline_args or sys.argv[1:]
 
@@ -98,7 +145,7 @@ def run(cmdline_args=None, set_retcode=True):
     event_aggregator = gokart.slack.EventAggregator()
     try:
         event_aggregator.set_handlers()
-        luigi.cmdline.luigi_run(cmdline_args)
+        _run_with_retcodes(cmdline_args)
     except SystemExit as e:
         _try_to_send_event_summary_to_slack(slack_api, event_aggregator, cmdline_args)
         sys.exit(e.code)

gokart/worker.py

@@ -67,9 +67,9 @@
 
 logger = logging.getLogger(__name__)
 
-# Set the start method to fork, which is the default on Unix systems.
-# This is necessary because the default start method on macOS is spawn, which is not compatible with the multiprocessing
-multiprocessing.set_start_method('fork')
+# Use fork context instead of the default (spawn on macOS), which ensures compatibility with gokart's multiprocessing requirements.
+_fork_context = multiprocessing.get_context('fork')
+_ForkProcess = _fork_context.Process
 
 # Prevent fork() from being called during a C-level getaddrinfo() which uses a process-global mutex,
 # that may not be unlocked in child process, resulting in the process being locked indefinitely.
@@ -106,7 +106,7 @@ def _get_retry_policy_dict(task: Task) -> dict[str, Any]:
 )
 
 
-class TaskProcess(multiprocessing.Process):
+class TaskProcess(_ForkProcess):  # type: ignore[valid-type, misc]
     """Wrap all task execution in this class.
 
     Mainly for convenience since this is run in a separate process."""
@@ -447,14 +447,14 @@ def __init__(
                 pass
 
         # Keep info about what tasks are running (could be in other processes)
-        self._task_result_queue: multiprocessing.Queue = multiprocessing.Queue()
+        self._task_result_queue: multiprocessing.Queue = _fork_context.Queue()
         self._running_tasks: dict[str, TaskProcess] = {}
         self._idle_since: datetime.datetime | None = None
 
         # mp-safe dictionary for caching completation checks across task processes
         self._task_completion_cache = None
         if self._config.cache_task_completion:
-            self._task_completion_cache = multiprocessing.Manager().dict()
+            self._task_completion_cache = _fork_context.Manager().dict()
 
         # Stuff for execution_summary
         self._add_task_history: list[Any] = []
@@ -641,8 +641,8 @@ def add(self, task: Task, multiprocess: bool = False, processes: int = 0) -> boo
             self._first_task = task.task_id
         self.add_succeeded = True
         if multiprocess:
-            queue: Any = multiprocessing.Manager().Queue()
-            pool: Any = multiprocessing.Pool(processes=processes if processes > 0 else None)
+            queue: Any = _fork_context.Manager().Queue()
+            pool: Any = _fork_context.Pool(processes=processes if processes > 0 else None)
         else:
             queue = luigi.worker.DequeQueue()
             pool = luigi.worker.SingleProcessPool()
@@ -905,7 +905,7 @@ def _run_task(self, task_id: str) -> None:
             task_process.run()
 
     def _create_task_process(self, task):
-        message_queue: Any = multiprocessing.Queue() if task.accepts_messages else None
+        message_queue: Any = _fork_context.Queue() if task.accepts_messages else None
         reporter = luigi.worker.TaskStatusReporter(self._scheduler, task.task_id, self._id, message_queue)
         use_multiprocessing = self._config.force_multiprocessing or bool(self.worker_processes > 1)
         return ContextManagedTaskProcess(