Force spawn start method

elk/extraction/extraction.py

@@ -270,6 +270,9 @@ def get_splits() -> SplitDict:
  )
  for (split_name, split_info) in get_splits().items()
  }
+ import multiprocess as mp
+
+ mp.set_start_method("spawn", force=True) # type: ignore[attr-defined]
 
  ds = dict()
  for split, builder in builders.items():

elk/run.py

@@ -135,7 +135,8 @@ def apply_to_layers(
  layers = self.concatenate(layers)
 
  # Should we write to different CSV files for elicit vs eval?
- with mp.Pool(num_devices) as pool, open(self.out_dir / "eval.csv", "w") as f:
+ ctx = mp.get_context("spawn")
+ with ctx.Pool(num_devices) as pool, open(self.out_dir / "eval.csv", "w") as f:
  mapper = pool.imap_unordered if num_devices > 1 else map
  iterator: Iterator[Log] = tqdm( # type: ignore
  mapper(func, layers), total=len(layers)

tests/test_write_iterator_to_file.py

@@ -156,7 +156,8 @@ def to_csv_line(x):
  return x.to_csv_line(skip_baseline=True)
 
  try:
- with mp.Pool(processes) as pool, open(tmp_path / "eval.csv", "w") as f:
+ ctx = mp.get_context("spawn")
+ with ctx.Pool(processes) as pool, open(tmp_path / "eval.csv", "w") as f:
  layers = [1, 2, 3]
  iterator = pool.imap_unordered(log_function, layers)
  write_iterator_to_file(