Fix benchmark_single_table with custom synthesizers and timeout (#337)

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
  'appdirs>=1.3',
  'boto3>=1.28,<2',
  'botocore>=1.31,<2',
+ 'cloudpickle>=2.1.0',
  'compress-pickle>=1.2.0',
  'humanfriendly>=8.2',
  "numpy>=1.21.0,<2.0.0;python_version<'3.10'",

sdgym/benchmark.py

@@ -7,10 +7,12 @@
 import pickle
 import tracemalloc
 import warnings
+from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
 
 import boto3
+import cloudpickle
 import compress_pickle
 import numpy as np
 import pandas as pd
@@ -318,6 +320,26 @@ def _score(
  return output
 
 
+@contextmanager
+def multiprocessing_context():
+ """Override multiprocessing ForkingPickler to use cloudpickle."""
+ original_dump = multiprocessing.reduction.ForkingPickler.dumps
+ original_load = multiprocessing.reduction.ForkingPickler.loads
+ original_method = multiprocessing.get_start_method()
+
+ multiprocessing.set_start_method('spawn', force=True)
+ multiprocessing.reduction.ForkingPickler.dumps = cloudpickle.dumps
+ multiprocessing.reduction.ForkingPickler.loads = cloudpickle.loads
+
+ try:
+ yield
+ finally:
+ # Restore original methods
+ multiprocessing.set_start_method(original_method, force=True)
+ multiprocessing.reduction.ForkingPickler.dumps = original_dump
+ multiprocessing.reduction.ForkingPickler.loads = original_load
+
+
 def _score_with_timeout(
  timeout,
  synthesizer,
@@ -329,32 +351,33 @@ def _score_with_timeout(
  modality=None,
  dataset_name=None,
 ):
- with multiprocessing.Manager() as manager:
- output = manager.dict()
- process = multiprocessing.Process(
- target=_score,
- args=(
- synthesizer,
- data,
- metadata,
- metrics,
- output,
- compute_quality_score,
- compute_diagnostic_score,
- modality,
- dataset_name,
- ),
- )
+ with multiprocessing_context():
+ with multiprocessing.Manager() as manager:
+ output = manager.dict()
+ process = multiprocessing.Process(
+ target=_score,
+ args=(
+ synthesizer,
+ data,
+ metadata,
+ metrics,
+ output,
+ compute_quality_score,
+ compute_diagnostic_score,
+ modality,
+ dataset_name,
+ ),
+ )
 
- process.start()
- process.join(timeout)
- process.terminate()
+  process.start()
+  process.join(timeout)
+  process.terminate()
 
- output = dict(output)
- if output.get('timeout'):
- LOGGER.error('Timeout running %s on dataset %s;', synthesizer['name'], dataset_name)
+  output = dict(output)
+  if output.get('timeout'):
+  LOGGER.error('Timeout running %s on dataset %s;', synthesizer['name'], dataset_name)
 
- return output
+  return output
 
 
 def _format_output(

sdgym/synthesizers/generate.py

@@ -124,7 +124,7 @@ def get_trained_synthesizer(self, data, metadata):
  obj:
  The trained synthesizer.
  """
- return get_trained_synthesizer_fn(data, metadata)
+ return self.synthesizer_fn['get_trained_synthesizer_fn'](data, metadata)
 
  def sample_from_synthesizer(self, synthesizer, num_samples):
  """Sample the desired number of samples from the given synthesizer.
@@ -139,11 +139,22 @@ def sample_from_synthesizer(self, synthesizer, num_samples):
  pandas.DataFrame:
  The synthetic data.
  """
- return sample_from_synthesizer_fn(synthesizer, num_samples)
-
- NewSynthesizer.__name__ = f'Custom:{display_name}'
-
- return NewSynthesizer
+ return self.synthesizer_fn['sample_from_synthesizer_fn'](synthesizer, num_samples)
+
+ CustomSynthesizer = type(
+ f'Custom:{display_name}',
+ (NewSynthesizer,),
+ {
+ 'synthesizer_fn': {
+ 'get_trained_synthesizer_fn': get_trained_synthesizer_fn,
+ 'sample_from_synthesizer_fn': sample_from_synthesizer_fn,
+ },
+ },
+ )
+ CustomSynthesizer.__name__ = f'Custom:{display_name}'
+ CustomSynthesizer.__module__ = 'sdgym.synthesizers.generate'
+ globals()[f'Custom:{display_name}'] = CustomSynthesizer
+ return CustomSynthesizer
 
 
 def create_multi_table_synthesizer(