improve datetime logic

Author: R-Palazzo

.github/workflows/run_benchmark.yml

@@ -9,7 +9,7 @@ on:
     - cron: '0 5 1 * *'
 
 jobs:
-  sdgym-benchmark:
+  run-sdgym-benchmark:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
@@ -31,4 +31,4 @@ jobs:
         AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
         AWS_DEFAULT_REGION: ${{ secrets.AWS_REGION }}
 
-      run: invoke sdgym-benchmark
+      run: invoke run-sdgym-benchmark

.github/workflows/upload_benchmark_results.yml

@@ -9,15 +9,11 @@ on:
     types:
       - completed
   workflow_dispatch:
-    inputs:
-      date:
-        description: 'Benchmark date (YYYY-MM-DD), defaults to the first of the current month'
-        required: false
   schedule:
     - cron: '0 6 * * *'
 
 jobs:
-  sdgym-benchmark:
+  upload-sdgym-benchmark:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
@@ -42,11 +38,4 @@ jobs:
         AWS_DEFAULT_REGION: ${{ secrets.AWS_REGION }}
 
       run: |
-        if [ -z "${{ github.event.inputs.date }}" ]; then
-          BENCHMARK_DATE=$(date -u "+%Y-%m-01")
-        else
-          BENCHMARK_DATE="${{ github.event.inputs.date }}"
-        fi
-
-        echo "Benchmark date: $BENCHMARK_DATE"
-        invoke upload-benchmark-results --date "$BENCHMARK_DATE"
+        invoke upload-benchmark-results

sdgym/_run_benchmark/_utils.py

@@ -0,0 +1,10 @@
+from datetime import datetime
+
+
+def get_run_name(date_str):
+    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

@@ -1,14 +1,41 @@
+import json
 import os
+from datetime import datetime, timezone
 
 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.s3 import get_s3_client, parse_s3_path
 
 datasets = ['expedia_hotel_logs', 'fake_companies']  # DEFAULT_DATASETS
 
 
+def append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str):
+    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'
+    try:
+        object = s3_client.get_object(Bucket=bucket, Key=f'{prefix}{key}')
+        body = object['Body'].read().decode('utf-8')
+        data = json.loads(body)
+    except s3_client.exceptions.ClientError as e:
+        if e.response['Error']['Code'] == 'NoSuchKey':
+            data = {'runs': []}
+        else:
+            raise RuntimeError(f'Failed to read {key} from S3: {e}')
+
+    data['runs'].append({'date': date_str, 'run_name': get_run_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'))
+
+
 def main():
     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 ['GaussianCopulaSynthesizer', 'TVAESynthesizer']:
         benchmark_single_table_aws(
             output_destination=run_benchmark.OUTPUT_DESTINATION_AWS,
@@ -19,6 +46,8 @@ def main():
             compute_privacy_score=False,
         )
 
+    append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str)
+
 
 if __name__ == '__main__':
     main()

sdgym/_run_benchmark/upload_benchmark_results.py

@@ -1,8 +1,7 @@
-import argparse
+import json
 import logging
 import os
 import sys
-from datetime import datetime
 
 import boto3
 from botocore.exceptions import ClientError
@@ -15,19 +14,15 @@
 LOGGER = logging.getLogger(__name__)
 
 
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--date', type=str, help='Benchmark date (YYYY-MM-DD)')
-    return parser.parse_args()
-
-
-def get_run_name(date_str):
+def get_latest_run_from_file(s3_client, bucket, key):
     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}'
+        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['run_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, run_name):
@@ -48,20 +43,14 @@ def upload_already_done(s3_client, bucket, prefix, run_name):
 
 
 def get_run_name_and_s3_vars(aws_access_key_id, aws_secret_access_key):
-    args = parse_args()
-    if args.date:
-        date_str = args.date
-    else:
-        date_str = datetime.utcnow().replace(day=1).strftime('%Y-%m-%d')
-
-    run_name = get_run_name(date_str)
     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,
     )
+    run_name = get_latest_run_from_file(s3_client, bucket, f'{prefix}_BENCHMARK_DATES.json')
 
     return run_name, s3_client, bucket, prefix
 
@@ -75,7 +64,7 @@ def upload_results(aws_access_key_id, aws_secret_access_key, run_name, s3_client
     result_writer = S3ResultsWriter(s3_client)
 
     if not result_explorer.all_runs_complete(run_name):
-        LOGGER.info(f'Run {run_name} is not complete yet. Exiting.')
+        LOGGER.warning(f'Run {run_name} is not complete yet. Exiting.')
         sys.exit(0)
 
     LOGGER.info(f'Run {run_name} is complete! Proceeding with summarization...')

tasks.py

@@ -204,7 +204,7 @@ def rmdir(c, path):
         pass
 
 @task
-def sdgym_benchmark(c):
+def run_sdgym_benchmark(c):
     """Run the SDGym benchmark."""
     c.run('python sdgym/_run_benchmark/run_benchmark.py')
 

tests/unit/_run_benchmark/test__utils.py

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

tests/unit/_run_benchmark/test_run_benchmark.py

@@ -1,3 +1,4 @@
+from datetime import datetime, timezone
 from unittest.mock import call, patch
 
 from sdgym._run_benchmark import OUTPUT_DESTINATION_AWS
@@ -6,10 +7,12 @@
 
 @patch('sdgym._run_benchmark.run_benchmark.benchmark_single_table_aws')
 @patch('sdgym._run_benchmark.run_benchmark.os.getenv')
-def test_main(mock_getenv, mock_benchmark_single_table_aws):
+@patch('sdgym._run_benchmark.run_benchmark.append_benchmark_run')
+def test_main(mock_append_benchmark_run, mock_getenv, mock_benchmark_single_table_aws):
     """Test the `main` method."""
     # Setup
     mock_getenv.side_effect = ['my_access_key', 'my_secret_key']
+    date = datetime.now(timezone.utc).strftime('%Y-%m-%d')
 
     # Run
     main()
@@ -31,3 +34,8 @@ def test_main(mock_getenv, mock_benchmark_single_table_aws):
         )
 
     mock_benchmark_single_table_aws.assert_has_calls(expected_calls)
+    mock_append_benchmark_run.assert_called_once_with(
+        'my_access_key',
+        'my_secret_key',
+        date,
+    )

tests/unit/_run_benchmark/test_upload_benchmark_result.py

@@ -4,47 +4,15 @@
 from botocore.exceptions import ClientError
 
 from sdgym._run_benchmark.upload_benchmark_results import (
-    get_run_name,
     get_run_name_and_s3_vars,
     main,
-    parse_args,
     upload_already_done,
     upload_results,
     write_uploaded_marker,
 )
 from sdgym.s3 import S3_REGION
 
 
-@patch('sdgym._run_benchmark.upload_benchmark_results.argparse.ArgumentParser')
-def test_parse_args(mock_argparse):
-    """Test the `parse_args` method."""
-    # Setup
-    parser = mock_argparse.return_value
-    parser.parse_args.return_value = Mock(date='01-07-2025')
-    mock_argparse.return_value.add_argument = Mock()
-
-    # Run
-    args = parse_args()
-
-    # Assert
-    assert args.date == '01-07-2025'
-    parser.add_argument.assert_called_once_with(
-        '--date', type=str, help='Benchmark date (YYYY-MM-DD)'
-    )
-    parser.parse_args.assert_called_once()
-
-
-def test_get_run_name():
-    """Test the `get_run_name` method."""
-    # Setup
-    expected_error_message = 'Invalid date format: invalid-date. Expected YYYY-MM-DD.'
-
-    # Run and Assert
-    assert get_run_name('2023-10-01') == 'SDGym_results_10_01_2023'
-    with pytest.raises(ValueError, match=expected_error_message):
-        get_run_name('invalid-date')
-
-
 def test_write_uploaded_marker():
     """Test the `write_uploaded_marker` method."""
     # Setup
@@ -92,41 +60,40 @@ def test_upload_already_done():
     assert result_false is False
 
 
-@patch('sdgym._run_benchmark.upload_benchmark_results.get_run_name')
 @patch('sdgym._run_benchmark.upload_benchmark_results.boto3.client')
 @patch('sdgym._run_benchmark.upload_benchmark_results.parse_s3_path')
 @patch('sdgym._run_benchmark.upload_benchmark_results.OUTPUT_DESTINATION_AWS')
-@patch('sdgym._run_benchmark.upload_benchmark_results.parse_args')
+@patch('sdgym._run_benchmark.upload_benchmark_results.get_latest_run_from_file')
 def test_get_run_name_and_s3_vars(
-    mock_parse_args,
+    mock_get_latest_run_from_file,
     mock_output_destination_aws,
     mock_parse_s3_path,
     mock_boto_client,
-    mock_get_run_name,
 ):
     """Test the `get_run_name_and_s3_vars` method."""
     # Setup
-    mock_parse_args.return_value.date = '2023-10-01'
     aws_access_key_id = 'my_access_key'
     aws_secret_access_key = 'my_secret_key'
     expected_result = ('SDGym_results_10_01_2023', 's3_client', 'bucket', 'prefix')
-    mock_get_run_name.return_value = 'SDGym_results_10_01_2023'
     mock_boto_client.return_value = 's3_client'
     mock_parse_s3_path.return_value = ('bucket', 'prefix')
+    mock_get_latest_run_from_file.return_value = 'SDGym_results_10_01_2023'
 
     # Run
     result = get_run_name_and_s3_vars(aws_access_key_id, aws_secret_access_key)
 
     # Assert
     assert result == expected_result
-    mock_get_run_name.assert_called_once_with('2023-10-01')
     mock_boto_client.assert_called_once_with(
         's3',
         aws_access_key_id=aws_access_key_id,
         aws_secret_access_key=aws_secret_access_key,
         region_name=S3_REGION,
     )
     mock_parse_s3_path.assert_called_once_with(mock_output_destination_aws)
+    mock_get_latest_run_from_file.assert_called_once_with(
+        's3_client', 'bucket', 'prefix_BENCHMARK_DATES.json'
+    )
 
 
 @patch('sdgym._run_benchmark.upload_benchmark_results.SDGymResultsExplorer')
@@ -202,7 +169,7 @@ def test_upload_results_not_all_runs_complete(
         )
 
     # Assert
-    mock_logger.info.assert_called_once_with(f'Run {run_name} is not complete yet. Exiting.')
+    mock_logger.warning.assert_called_once_with(f'Run {run_name} is not complete yet. Exiting.')
     mock_sdgym_results_explorer.assert_called_once_with(
         mock_output_destination_aws,
         aws_access_key_id=aws_access_key_id,