address comments

Author: R-Palazzo

.github/workflows/run_benchmark.yml

@@ -22,7 +22,6 @@ jobs:
     - name: Install dependencies
       run: |
           python -m pip install --upgrade pip
-          python -m pip install invoke
           python -m pip install -e .[dev]
           
     - name: Run SDGym Benchmark

sdgym/_run_benchmark/__init__.py

@@ -1,44 +1,53 @@
+"""Script to run a benchmark and upload results to S3."""
+
 import json
 import os
 from datetime import datetime, timezone
 
 from botocore.exceptions import ClientError
 
-import sdgym._run_benchmark as run_benchmark
-from sdgym._run_benchmark._utils import get_run_name
 from sdgym.benchmark import benchmark_single_table_aws
+from sdgym.run_benchmark.utils import (
+    KEY_DATE_FILE,
+    OUTPUT_DESTINATION_AWS,
+    SYNTHESIZERS,
+    get_result_folder_name,
+)
 from sdgym.s3 import get_s3_client, parse_s3_path
 
 
 def append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str):
+    """Append a new benchmark run to the benchmark dates file in S3."""
     s3_client = get_s3_client(
         aws_access_key_id=aws_access_key_id,
         aws_secret_access_key=aws_secret_access_key,
     )
-    bucket, prefix = parse_s3_path(run_benchmark.OUTPUT_DESTINATION_AWS)
-    key = '_BENCHMARK_DATES.json'
+    bucket, prefix = parse_s3_path(OUTPUT_DESTINATION_AWS)
     try:
-        object = s3_client.get_object(Bucket=bucket, Key=f'{prefix}{key}')
+        object = s3_client.get_object(Bucket=bucket, Key=f'{prefix}{KEY_DATE_FILE}')
         body = object['Body'].read().decode('utf-8')
         data = json.loads(body)
     except ClientError as e:
         if e.response['Error']['Code'] == 'NoSuchKey':
             data = {'runs': []}
         else:
-            raise RuntimeError(f'Failed to read {key} from S3: {e}')
+            raise RuntimeError(f'Failed to read {KEY_DATE_FILE} from S3: {e}')
 
-    data['runs'].append({'date': date_str, 'run_name': get_run_name(date_str)})
+    data['runs'].append({'date': date_str, 'folder_name': get_result_folder_name(date_str)})
     data['runs'] = sorted(data['runs'], key=lambda x: x['date'])
-    s3_client.put_object(Bucket=bucket, Key=f'{prefix}{key}', Body=json.dumps(data).encode('utf-8'))
+    s3_client.put_object(
+        Bucket=bucket, Key=f'{prefix}{KEY_DATE_FILE}', Body=json.dumps(data).encode('utf-8')
+    )
 
 
 def main():
+    """Main function to run the benchmark and upload results."""
     aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
     aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
     date_str = datetime.now(timezone.utc).strftime('%Y-%m-%d')
-    for synthesizer in run_benchmark.SYNTHESIZERS:
+    for synthesizer in SYNTHESIZERS:
         benchmark_single_table_aws(
-            output_destination=run_benchmark.OUTPUT_DESTINATION_AWS,
+            output_destination=OUTPUT_DESTINATION_AWS,
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
             synthesizers=[synthesizer],

sdgym/_run_benchmark/_utils.py

@@ -0,0 +1,102 @@
+"""Script to upload benchmark results to S3."""
+
+import json
+import logging
+import os
+import sys
+
+import boto3
+from botocore.exceptions import ClientError
+
+from sdgym.result_writer import S3ResultsWriter
+from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS
+from sdgym.s3 import S3_REGION, parse_s3_path
+from sdgym.sdgym_result_explorer.result_explorer import SDGymResultsExplorer
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_latest_run_from_file(s3_client, bucket, key):
+    """Get the latest run folder name from the benchmark dates file in S3."""
+    try:
+        object = s3_client.get_object(Bucket=bucket, Key=key)
+        body = object['Body'].read().decode('utf-8')
+        data = json.loads(body)
+        latest = sorted(data['runs'], key=lambda x: x['date'])[-1]
+        return latest['folder_name']
+    except s3_client.exceptions.ClientError as e:
+        raise RuntimeError(f'Failed to read {key} from S3: {e}')
+
+
+def write_uploaded_marker(s3_client, bucket, prefix, folder_name):
+    """Write a marker file to indicate that the upload is complete."""
+    s3_client.put_object(
+        Bucket=bucket, Key=f'{prefix}{folder_name}/upload_complete.marker', Body=b'Upload complete'
+    )
+
+
+def upload_already_done(s3_client, bucket, prefix, folder_name):
+    """Check if the upload has already been done by looking for the marker file."""
+    try:
+        s3_client.head_object(Bucket=bucket, Key=f'{prefix}{folder_name}/upload_complete.marker')
+        return True
+    except ClientError as e:
+        if e.response['Error']['Code'] == '404':
+            return False
+
+        raise
+
+
+def get_result_folder_name_and_s3_vars(aws_access_key_id, aws_secret_access_key):
+    """Get the result folder name and S3 client variables."""
+    bucket, prefix = parse_s3_path(OUTPUT_DESTINATION_AWS)
+    s3_client = boto3.client(
+        's3',
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
+        region_name=S3_REGION,
+    )
+    folder_name = get_latest_run_from_file(s3_client, bucket, f'{prefix}_BENCHMARK_DATES.json')
+
+    return folder_name, s3_client, bucket, prefix
+
+
+def upload_results(
+    aws_access_key_id, aws_secret_access_key, folder_name, s3_client, bucket, prefix
+):
+    """Upload benchmark results to S3."""
+    result_explorer = SDGymResultsExplorer(
+        OUTPUT_DESTINATION_AWS,
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
+    )
+    result_writer = S3ResultsWriter(s3_client)
+
+    if not result_explorer.all_runs_complete(folder_name):
+        LOGGER.warning(f'Run {folder_name} is not complete yet. Exiting.')
+        sys.exit(0)
+
+    LOGGER.info(f'Run {folder_name} is complete! Proceeding with summarization...')
+    summary, _ = result_explorer.summarize(folder_name)
+    result_writer.write_dataframe(
+        summary, f'{OUTPUT_DESTINATION_AWS}{folder_name}/{folder_name}_summary.csv', index=True
+    )
+    write_uploaded_marker(s3_client, bucket, prefix, folder_name)
+
+
+def main():
+    """Main function to upload benchmark results."""
+    aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
+    aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
+    folder_name, s3_client, bucket, prefix = get_result_folder_name_and_s3_vars(
+        aws_access_key_id, aws_secret_access_key
+    )
+    if upload_already_done(s3_client, bucket, prefix, folder_name):
+        LOGGER.warning('Benchmark results have already been uploaded. Exiting.')
+        sys.exit(0)
+
+    upload_results(aws_access_key_id, aws_secret_access_key, folder_name, s3_client, bucket, prefix)
+
+
+if __name__ == '__main__':
+    main()

sdgym/_run_benchmark/upload_benchmark_results.py

@@ -0,0 +1,22 @@
+"""Utils file for the run_benchmark module."""
+
+from datetime import datetime
+
+from sdgym.benchmark import SDV_SINGLE_TABLE_SYNTHESIZERS
+
+OUTPUT_DESTINATION_AWS = 's3://sdgym-benchmark/Debug/Issue_425/'
+UPLOAD_DESTINATION_AWS = 's3://sdgym-benchmark/Debug/Issue_425/'
+DEBUG_SLACK_CHANNEL = 'sdv-alerts-debug'
+SLACK_CHANNEL = 'sdv-alerts'
+KEY_DATE_FILE = '_BENCHMARK_DATES.json'
+SYNTHESIZERS = SDV_SINGLE_TABLE_SYNTHESIZERS
+
+
+def get_result_folder_name(date_str):
+    """Get the result folder name based on the date string."""
+    try:
+        date = datetime.strptime(date_str, '%Y-%m-%d')
+    except ValueError:
+        raise ValueError(f'Invalid date format: {date_str}. Expected YYYY-MM-DD.')
+
+    return f'SDGym_results_{date.month:02d}_{date.day:02d}_{date.year}'

sdgym/_run_benchmark/run_benchmark.py sdgym/run_benchmark/run_benchmark.pysdgym/_run_benchmark/run_benchmark.py renamed to sdgym/run_benchmark/run_benchmark.py

@@ -0,0 +1,14 @@
+import pytest
+
+from sdgym.run_benchmark.utils import get_result_folder_name
+
+
+def test_get_result_folder_name():
+    """Test the `get_result_folder_name` method."""
+    # Setup
+    expected_error_message = 'Invalid date format: invalid-date. Expected YYYY-MM-DD.'
+
+    # Run and Assert
+    assert get_result_folder_name('2023-10-01') == 'SDGym_results_10_01_2023'
+    with pytest.raises(ValueError, match=expected_error_message):
+        get_result_folder_name('invalid-date')