add unit tests

R-Palazzo · R-Palazzo · commit 3eed4feed98c · 2025-08-04T14:37:32.000+01:00
diff --git a/.github/workflows/upload_benchmark_results.yml b/.github/workflows/upload_benchmark_results.yml
@@ -47,43 +47,48 @@ jobs:
           echo "Upload skipped. Exiting workflow."
           exit 0
         fi
-    - name: Prepare summary file for PR
+    - name: Prepare files for PR
       run: |
         mkdir pr-staging
-        echo "Looking for: $GITHUB_LOCAL_RESULTS_DIR/${FOLDER_NAME}_summary.csv"
+        echo "Looking for files in: $GITHUB_LOCAL_RESULTS_DIR"
         ls -l "$GITHUB_LOCAL_RESULTS_DIR"
-        cp "$GITHUB_LOCAL_RESULTS_DIR/${FOLDER_NAME}_summary.csv" \
-           "pr-staging/SDGym_summary_${FOLDER_NAME}.csv"
+        for f in "$GITHUB_LOCAL_RESULTS_DIR"/${FOLDER_NAME}_*.csv; do
+          base=$(basename "$f")
+          cp "$f" "pr-staging/SDGym_${base}"
+        done
+
+        echo "Files staged for PR:"
+        ls -l pr-staging
 
     - name: Checkout target repo (sdv-dev.github.io)
       run: |
         git clone https://github.com/sdv-dev/sdv-dev.github.io.git target-repo
         cd target-repo
         git checkout gatsby-home
-    - name: Copy summary and create PR
+
+    - name: Copy results and create PR
       env:
+        GH_TOKEN: ${{ secrets.GH_TOKEN }}
         FOLDER_NAME: ${{ env.FOLDER_NAME }}
       run: |
-        cp pr-staging/SDGym_summary_${FOLDER_NAME}.csv target-repo/assets/
-
+        cp pr-staging/* target-repo/assets/
         cd target-repo
         git checkout -b sdgym-benchmark-upload-${FOLDER_NAME}
-        git config --local user.name "${GITHUB_ACTOR}"
-        git config --local user.email "${GITHUB_ACTOR_ID}+${GITHUB_ACTOR}@users.noreply.github.com"
+        git config --local user.name "github-actions[bot]"
+        git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com"
 
         git add assets/
-        git commit -m "Upload SDGym Benchmark Summary ($FOLDER_NAME)"
-
-        git remote set-url origin https://x-access-token:${{ secrets.GITHUB_TOKEN }}@github.com/sdv-dev/sdv-dev.github.io.git
+        git commit -m "Upload SDGym Benchmark Results ($FOLDER_NAME)"
+        git remote set-url origin https://x-access-token:${GH_TOKEN}@github.com/sdv-dev/sdv-dev.github.io.git
         git push origin sdgym-benchmark-upload-${FOLDER_NAME}
 
+        # Create PR
         gh pr create \
           --repo sdv-dev/sdv-dev.github.io \
           --head sdgym-benchmark-upload-${FOLDER_NAME} \
           --base gatsby-home \
-          --title "Upload SDGym Benchmark Summary ($FOLDER_NAME)" \
-          --body "Automated benchmark summary upload" \
-          --assignee "${{ github.actor }}"
+          --title "Upload SDGym Benchmark Results ($FOLDER_NAME)" \
+          --body "Automated SDGym benchmark results upload"
 
     - name: Send Slack notification
       env:
diff --git a/sdgym/run_benchmark/upload_benchmark_results.py b/sdgym/run_benchmark/upload_benchmark_results.py
@@ -9,7 +9,7 @@
 from botocore.exceptions import ClientError
 
 from sdgym.result_writer import LocalResultsWriter, S3ResultsWriter
-from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS
+from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS, get_df_to_plot
 from sdgym.s3 import S3_REGION, parse_s3_path
 from sdgym.sdgym_result_explorer.result_explorer import SDGymResultsExplorer
 
@@ -86,7 +86,8 @@ def upload_results(
                 env_file.write('SKIP_UPLOAD=false\n')
                 env_file.write(f'FOLDER_NAME={folder_name}\n')
 
-    summary, _ = result_explorer.summarize(folder_name)
+    summary, results = result_explorer.summarize(folder_name)
+    df_to_plot = get_df_to_plot(results)
     result_writer.write_dataframe(
         summary, f'{OUTPUT_DESTINATION_AWS}{folder_name}/{folder_name}_summary.csv', index=True
     )
@@ -96,6 +97,9 @@ def upload_results(
         local_results_writer.write_dataframe(
             summary, f'{local_export_dir}/{folder_name}_summary.csv', index=True
         )
+        local_results_writer.write_dataframe(
+            df_to_plot, f'{local_export_dir}/{folder_name}_plot_data.csv', index=False
+        )
 
     write_uploaded_marker(s3_client, bucket, prefix, folder_name)
 
diff --git a/sdgym/run_benchmark/utils.py b/sdgym/run_benchmark/utils.py
@@ -3,6 +3,7 @@
 import os
 from datetime import datetime
 
+import numpy as np
 from slack_sdk import WebClient
 
 from sdgym.s3 import parse_s3_path
@@ -12,6 +13,37 @@
 DEBUG_SLACK_CHANNEL = 'sdv-alerts-debug'
 SLACK_CHANNEL = 'sdv-alerts'
 KEY_DATE_FILE = '_BENCHMARK_DATES.json'
+PLOTLY_MARKERS = [
+    'circle',
+    'square',
+    'diamond',
+    'cross',
+    'x',
+    'triangle-up',
+    'triangle-down',
+    'triangle-left',
+    'triangle-right',
+    'pentagon',
+    'hexagon',
+    'hexagon2',
+    'octagon',
+    'star',
+    'hexagram',
+    'star-triangle-up',
+    'star-triangle-down',
+    'star-square',
+    'star-diamond',
+    'diamond-tall',
+    'diamond-wide',
+    'hourglass',
+    'bowtie',
+    'circle-cross',
+    'circle-x',
+    'square-cross',
+    'square-x',
+    'diamond-cross',
+    'diamond-x',
+]
 
 # The synthesizers inside the same list will be run by the same ec2 instance
 SYNTHESIZERS_SPLIT = [
@@ -68,7 +100,7 @@ def post_benchmark_launch_message(date_str):
     bucket, prefix = parse_s3_path(OUTPUT_DESTINATION_AWS)
     url_link = get_s3_console_link(bucket, f'{prefix}{folder_name}/')
     body = '🏃 SDGym benchmark has been launched! EC2 Instances are running. '
-    body += f'Intermediate results can be found <{url_link} |here>.\n'
+    body += f'Intermediate results can be found <{url_link}|here>.\n'
     post_slack_message(channel, body)
 
 
@@ -85,3 +117,43 @@ def post_benchmark_uploaded_message(folder_name, pr_url=None):
         body += f'Waiting on merging this PR to update GitHub directory: <{pr_url}|PR Link>\n'
 
     post_slack_message(channel, body)
+
+
+def get_df_to_plot(benchmark_result):
+    """Get the data to plot from the benchmark result.
+
+    Args:
+        benchmark_result (DataFrame): The benchmark result DataFrame.
+
+    Returns:
+        DataFrame: The data to plot.
+    """
+    df_to_plot = benchmark_result.copy()
+    df_to_plot['total_time'] = df_to_plot['Train_Time'] + df_to_plot['Sample_Time']
+    df_to_plot['Aggregated_Time'] = df_to_plot.groupby('Synthesizer')['total_time'].transform('sum')
+    df_to_plot = (
+        df_to_plot.groupby('Synthesizer')[['Aggregated_Time', 'Quality_Score']].mean().reset_index()
+    )
+    df_to_plot['Log10 Aggregated_Time'] = df_to_plot['Aggregated_Time'].apply(
+        lambda x: np.log10(x) if x > 0 else 0
+    )
+    df_to_plot = df_to_plot.sort_values(
+        ['Aggregated_Time', 'Quality_Score'], ascending=[True, False]
+    )
+    df_to_plot['Cumulative Quality Score'] = df_to_plot['Quality_Score'].cummax()
+    pareto_points = df_to_plot.loc[
+        df_to_plot['Quality_Score'] == df_to_plot['Cumulative Quality Score']
+    ]
+    df_to_plot['Pareto'] = df_to_plot.index.isin(pareto_points.index)
+    df_to_plot['Color'] = df_to_plot['Pareto'].apply(lambda x: '#01E0C9' if x else '#03AFF1')
+    df_to_plot['Synthesizer'] = df_to_plot['Synthesizer'].str.replace(
+        'Synthesizer', '', regex=False
+    )
+
+    synthesizers = df_to_plot['Synthesizer'].unique()
+    marker_map = {
+        synth: PLOTLY_MARKERS[i % len(PLOTLY_MARKERS)] for i, synth in enumerate(synthesizers)
+    }
+    df_to_plot['Marker'] = df_to_plot['Synthesizer'].map(marker_map)
+
+    return df_to_plot.drop(columns=['Cumulative Quality Score']).reset_index(drop=True)
diff --git a/tests/unit/run_benchmark/test__utils.py b/tests/unit/run_benchmark/test__utils.py
diff --git a/tests/unit/run_benchmark/test_upload_benchmark_result.py b/tests/unit/run_benchmark/test_upload_benchmark_result.py
@@ -1,4 +1,4 @@
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
 
 import pytest
 from botocore.exceptions import ClientError
@@ -103,7 +103,9 @@ def test_get_result_folder_name_and_s3_vars(
 @patch('sdgym.run_benchmark.upload_benchmark_results.OUTPUT_DESTINATION_AWS')
 @patch('sdgym.run_benchmark.upload_benchmark_results.LocalResultsWriter')
 @patch('sdgym.run_benchmark.upload_benchmark_results.os.environ.get')
+@patch('sdgym.run_benchmark.upload_benchmark_results.get_df_to_plot')
 def test_upload_results(
+    mock_get_df_to_plot,
     mock_os_environ_get,
     mock_local_results_writer,
     mock_output_destination_aws,
@@ -124,6 +126,7 @@ def test_upload_results(
     result_explorer_instance.all_runs_complete.return_value = True
     result_explorer_instance.summarize.return_value = ('summary', 'results')
     mock_os_environ_get.return_value = '/tmp/sdgym_results'
+    mock_get_df_to_plot.return_value = 'df_to_plot'
 
     # Run
     upload_results(
@@ -149,9 +152,13 @@ def test_upload_results(
     result_explorer_instance.summarize.assert_called_once_with(run_name)
     mock_s3_results_writer.return_value.write_dataframe.assert_called_once()
     mock_write_uploaded_marker.assert_called_once_with(s3_client, bucket, prefix, run_name)
-    mock_local_results_writer.return_value.write_dataframe.assert_called_once_with(
-        'summary', '/tmp/sdgym_results/SDGym_results_10_01_2023_summary.csv', index=True
-    )
+    mock_local_results_writer.return_value.write_dataframe.assert_has_calls([
+        call('summary', '/tmp/sdgym_results/SDGym_results_10_01_2023_summary.csv', index=True),
+        call(
+            'df_to_plot', '/tmp/sdgym_results/SDGym_results_10_01_2023_plot_data.csv', index=False
+        ),
+    ])
+    mock_get_df_to_plot.assert_called_once_with('results')
 
 
 @patch('sdgym.run_benchmark.upload_benchmark_results.SDGymResultsExplorer')
diff --git a/tests/unit/run_benchmark/test_utils.py b/tests/unit/run_benchmark/test_utils.py
