Force multiprocessing start method to 'fork' on macOS.

model_eval/cli/main.py

@@ -16,10 +16,21 @@
 
 import importlib
 import json
+import os
 import pathlib
 import sys
 from typing import Any
 
+if sys.platform == "darwin":
+  import multiprocessing  # pylint: disable=unused-import,g-import-not-at-top
+  # Disable macOS fork safety check to prevent crashes when calling Metal/GPU
+  # acceleration libraries.
+  os.environ["OBJC_DISABLE_INITIALIZE_FORK_SAFETY"] = "YES"
+  # Force multiprocessing start method to 'fork' to ensure compatibility with
+  # lighteval's nested functions. This is not ideal as fork() may be unsafe
+  # when combining with other multithreaded libraries.
+  multiprocessing.set_start_method("fork", force=True)
+
 from absl import app
 from absl import flags
 from model_eval.api import pipeline as eval_pipeline

model_eval/tests/unit/main_test.py

@@ -160,6 +160,21 @@ def test_list_runners_with_custom_config(self):
       self.assertIn("Supported runners for framework 'lm-eval':", result.output)
       self.assertIn("- my_custom_runner", result.output)
 
+  @mock.patch("sys.platform", "darwin")
+  @mock.patch("multiprocessing.set_start_method")
+  def test_multiprocessing_set_start_method_fork_on_mac(self, mock_set_start_method):
+    import importlib
+    # Reload main module to run module-level initialization code
+    importlib.reload(main)
+    mock_set_start_method.assert_called_once_with("fork", force=True)
+
+  @mock.patch("sys.platform", "linux")
+  @mock.patch("multiprocessing.set_start_method")
+  def test_multiprocessing_set_start_method_default_on_non_mac(self, mock_set_start_method):
+    import importlib
+    importlib.reload(main)
+    mock_set_start_method.assert_not_called()
+
 
 if __name__ == "__main__":
   unittest.main()