cleaning

Author: R-Palazzo

sdgym/benchmark.py

@@ -107,7 +107,7 @@
     'TVAESynthesizer',
 ]
 SDV_MULTI_TABLE_SYNTHESIZERS = ['HMASynthesizer']
-
+MODALITY_IDX = 10
 SDV_SYNTHESIZERS = SDV_SINGLE_TABLE_SYNTHESIZERS + SDV_MULTI_TABLE_SYNTHESIZERS
 
 
@@ -1486,7 +1486,7 @@ def _get_s3_script_content(
     return f"""
 import boto3
 import cloudpickle
-from sdgym.benchmark import _run_jobs, _write_metainfo_file, _update_metainfo_file
+from sdgym.benchmark import _run_jobs, _write_metainfo_file, _update_metainfo_file, MODALITY_IDX
 from io import StringIO
 from sdgym.result_writer import S3ResultsWriter
 
@@ -1498,8 +1498,9 @@ def _get_s3_script_content(
 )
 response = s3_client.get_object(Bucket='{bucket_name}', Key='{job_args_key}')
 job_args_list = cloudpickle.loads(response['Body'].read())
+modality = job_args_list[0][MODALITY_IDX]
 result_writer = S3ResultsWriter(s3_client=s3_client)
-_write_metainfo_file({synthesizers}, job_args_list, 'single_table', result_writer)
+_write_metainfo_file({synthesizers}, job_args_list, modality, result_writer)
 scores = _run_jobs(None, job_args_list, False, result_writer=result_writer)
 metainfo_filename = job_args_list[0][-1]['metainfo']
 _update_metainfo_file(metainfo_filename, result_writer)
@@ -1876,15 +1877,13 @@ def benchmark_multi_table_aws(
             Whether or not to evaluate an overall quality score. Defaults to ``True``.
         compute_diagnostic_score (bool):
             Whether or not to evaluate an overall diagnostic score. Defaults to ``True``.
-        compute_privacy_score (bool):
-            Whether or not to evaluate an overall privacy score. Defaults to ``True``.
         timeout (int or ``None``):
             The maximum number of seconds to wait for synthetic data creation. If ``None``, no
             timeout is enforced.
 
     Returns:
         pandas.DataFrame:
-            A table containing one row per synthesizer + dataset + metric.
+            A table containing one row per synthesizer + dataset.
     """
     s3_client = _validate_output_destination(
         output_destination,
@@ -1901,17 +1900,17 @@ def benchmark_multi_table_aws(
         limit_dataset_size=limit_dataset_size,
         sdv_datasets=sdv_datasets,
         additional_datasets_folder=additional_datasets_folder,
+        sdmetrics=None,
         timeout=timeout,
         output_destination=output_destination,
         compute_quality_score=compute_quality_score,
         compute_diagnostic_score=compute_diagnostic_score,
+        compute_privacy_score=None,
         synthesizers=synthesizers,
         detailed_results_folder=None,
         custom_synthesizers=None,
         s3_client=s3_client,
         modality='multi_table',
-        sdmetrics=None,
-        compute_privacy_score=None,
     )
     if not job_args_list:
         return _get_empty_dataframe(

sdgym/synthesizers/sdv.py

@@ -21,8 +21,6 @@
 
 MODEL_KWARGS = {'HMASynthesizer': {'verbose': False}}
 
-MODEL_KWARGS = {'HMASynthesizer': {'verbose': False}}
-
 
 def _get_sdv_synthesizers(modality):
     _validate_modality(modality)

tests/unit/test_benchmark.py

@@ -29,6 +29,7 @@
     _validate_output_destination,
     _write_metainfo_file,
     benchmark_multi_table,
+    benchmark_multi_table_aws,
     benchmark_single_table,
     benchmark_single_table_aws,
 )
@@ -1298,3 +1299,63 @@ def test_benchmark_multi_table_no_jobs(
     )
     mock__write_metainfo_file.assert_called_once()
     pd.testing.assert_frame_equal(scores, empty_scores)
+
+
+@patch('sdgym.benchmark._validate_output_destination')
+@patch('sdgym.benchmark._generate_job_args_list')
+@patch('sdgym.benchmark._run_on_aws')
+def test_benchmark_multi_table_aws(
+    mock_run_on_aws, mock_generate_job_args_list, mock_validate_output_destination
+):
+    """Test `benchmark_multi_table_aws` method."""
+    # Setup
+    output_destination = 's3://sdgym-benchmark/Debug/Issue_487_test_1'
+    synthesizers = ['HMASynthesizer']
+    datasets = ['financial', 'NBA']
+    aws_access_key_id = '12345'
+    aws_secret_access_key = '67890'
+    mock_validate_output_destination.return_value = 's3_client_mock'
+    mock_generate_job_args_list.return_value = 'job_args_list_mock'
+
+    # Run
+    benchmark_multi_table_aws(
+        output_destination=output_destination,
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
+        synthesizers=synthesizers,
+        sdv_datasets=datasets,
+    )
+
+    # Assert
+    assert 'MultiTableUniformSynthesizer' in synthesizers
+    mock_validate_output_destination.assert_called_once_with(
+        output_destination,
+        aws_keys={
+            'aws_access_key_id': aws_access_key_id,
+            'aws_secret_access_key': aws_secret_access_key,
+        },
+    )
+    mock_generate_job_args_list.assert_called_once_with(
+        limit_dataset_size=False,
+        sdv_datasets=datasets,
+        additional_datasets_folder=None,
+        sdmetrics=None,
+        timeout=None,
+        output_destination=output_destination,
+        compute_quality_score=True,
+        compute_diagnostic_score=True,
+        compute_privacy_score=None,
+        synthesizers=synthesizers,
+        detailed_results_folder=None,
+        custom_synthesizers=None,
+        s3_client='s3_client_mock',
+        modality='multi_table',
+    )
+    mock_run_on_aws.assert_called_once_with(
+        output_destination=output_destination,
+        synthesizers=synthesizers,
+        s3_client='s3_client_mock',
+        job_args_list='job_args_list_mock',
+        aws_access_key_id='12345',
+        aws_secret_access_key='67890',
+    )