update run_benchmark

R-Palazzo · R-Palazzo · commit 051dfa51bffe · 2025-12-12T17:45:13.000Z
diff --git a/.github/workflows/run_benchmark_multi_table.yml b/.github/workflows/run_benchmark_multi_table.yml
@@ -0,0 +1,35 @@
+name: Run SDGym Benchmark
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: '0 5 1 * *'
+
+jobs:
+  run-sdgym-benchmark:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Set up latest Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: 'pyproject.toml'
+    - name: Install dependencies
+      env:
+        username: ${{ secrets.GCP_USERNAME }}
+        license_key: ${{ secrets.GCP_LICENSE_KEY }}
+      run: |
+          python -m pip install --upgrade pip
+          python -m pip install bundle-xsynthesizers --index-url https://{username}:{license_key}@pypi.datacebo.com
+          python -m pip install --no-cache-dir -e .[dev]
+          
+    - name: Run SDGym Benchmark
+      env:
+        SLACK_TOKEN: ${{ secrets.SLACK_TOKEN }}
+        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        AWS_DEFAULT_REGION: ${{ secrets.AWS_REGION }}
+
+      run: invoke run-sdgym-benchmark --modality multi_table
diff --git a/sdgym/benchmark.py b/sdgym/benchmark.py
@@ -1654,31 +1654,38 @@ def _get_user_data_script(credentials, script_content, compute_service='aws'):
     else:
         termination_body = """
         echo "Terminating GCP instance via Compute API"
+        PROJECT_ID="$GCP_PROJECT"
+        ZONE="$GCP_ZONE"
+        INSTANCE="$INSTANCE_NAME"
 
-        TOKEN=$(curl -s -H "Metadata-Flavor: Google" \
+        TOKEN=$(curl -sf -H "Metadata-Flavor: Google" \
             http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token \
             | jq -r ".access_token" || true)
 
-        PROJECT_ID=$(curl -s -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/project/project-id || true)
+        echo "GCP termination parameters:"
+        echo "  PROJECT_ID=$PROJECT_ID"
+        echo "  ZONE=$ZONE"
+        echo "  INSTANCE=$INSTANCE"
 
-        ZONE=$(curl -s -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/instance/zone | awk -F/ '{print $4}' || true)
-
-        if [ -n "$TOKEN" ] && [ -n "$PROJECT_ID" ] && [ -n "$ZONE" ]; then
+        if [ -z "$PROJECT_ID" ] || [ -z "$ZONE" ] || [ -z "$INSTANCE" ] || [ -z "$TOKEN" ]; then
+            echo "Skipping GCP termination (missing required parameters)"
+        else
             curl -s -X DELETE \
                 -H "Authorization: Bearer $TOKEN" \
                 "https://compute.googleapis.com/compute/v1/projects/" \
                 "$PROJECT_ID/zones/$ZONE/instances/" \
-                "$HOSTNAME" \
+                "$INSTANCE" \
                 || true
         fi
-        """
-
+                """
     return textwrap.dedent(f"""\
         #!/bin/bash
         set -e
 
+        export GCP_PROJECT='{credentials['gcp']['gcp_project']}'
+        export GCP_ZONE='{credentials['gcp']['gcp_zone']}'
+        export INSTANCE_NAME='sdgym-run'
+
         terminate_instance() {{
         {termination_body}
         }}
diff --git a/sdgym/run_benchmark/run_benchmark.py b/sdgym/run_benchmark/run_benchmark.py
@@ -1,31 +1,33 @@
 """Script to run a benchmark and upload results to S3."""
 
+import argparse
 import json
 import os
 from datetime import datetime, timezone
 
 from botocore.exceptions import ClientError
 
-from sdgym.benchmark import benchmark_single_table_aws
+from sdgym.benchmark import _benchmark_multi_table_compute_gcp, benchmark_single_table_aws
 from sdgym.run_benchmark.utils import (
     KEY_DATE_FILE,
     OUTPUT_DESTINATION_AWS,
-    SYNTHESIZERS_SPLIT,
+    SYNTHESIZERS_SPLIT_MULTI_TABLE,
+    SYNTHESIZERS_SPLIT_SINGLE_TABLE,
     get_result_folder_name,
     post_benchmark_launch_message,
 )
 from sdgym.s3 import get_s3_client, parse_s3_path
 
 
-def append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str):
+def append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str, modality='single_table'):
     """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(OUTPUT_DESTINATION_AWS)
     try:
-        object = s3_client.get_object(Bucket=bucket, Key=f'{prefix}{KEY_DATE_FILE}')
+        object = s3_client.get_object(Bucket=bucket, Key=f'{prefix}{modality}{KEY_DATE_FILE}')
         body = object['Body'].read().decode('utf-8')
         data = json.loads(body)
     except ClientError as e:
@@ -41,23 +43,50 @@ def append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str):
     )
 
 
+def _parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--modality',
+        choices=['single_table', 'multi_table'],
+        default='single_table',
+        help='Benchmark modality to run.',
+    )
+    return parser.parse_args()
+
+
 def main():
     """Main function to run the benchmark and upload results."""
+    args = _parse_args()
     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_group in SYNTHESIZERS_SPLIT:
-        benchmark_single_table_aws(
-            output_destination=OUTPUT_DESTINATION_AWS,
-            aws_access_key_id=aws_access_key_id,
-            aws_secret_access_key=aws_secret_access_key,
-            synthesizers=synthesizer_group,
-            compute_privacy_score=False,
-            timeout=345600,  # 4 days
-        )
-
-    append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str)
-    post_benchmark_launch_message(date_str)
+
+    if args.modality == 'single_table':
+        for synthesizer_group in SYNTHESIZERS_SPLIT_SINGLE_TABLE:
+            benchmark_single_table_aws(
+                output_destination=OUTPUT_DESTINATION_AWS,
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+                synthesizers=synthesizer_group,
+                compute_privacy_score=False,
+                timeout=345600,  # 4 days
+            )
+
+        append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str, modality='single_table')
+
+    else:
+        for synthesizer_group in SYNTHESIZERS_SPLIT_MULTI_TABLE:
+            _benchmark_multi_table_compute_gcp(
+                output_destination='s3://sdgym-benchmark/Debug/GCP/',
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+                synthesizers=synthesizer_group,
+                compute_privacy_score=False,
+                timeout=345600,  # 4 days
+            )
+        append_benchmark_run(aws_access_key_id, aws_secret_access_key, date_str, modality='multi_table')
+
+    post_benchmark_launch_message(date_str, compute_service='GCP')
 
 
 if __name__ == '__main__':
diff --git a/sdgym/run_benchmark/utils.py b/sdgym/run_benchmark/utils.py
@@ -48,12 +48,16 @@
 ]
 
 # The synthesizers inside the same list will be run by the same ec2 instance
-SYNTHESIZERS_SPLIT = [
+SYNTHESIZERS_SPLIT_SINGLE_TABLE = [
     ['UniformSynthesizer', 'ColumnSynthesizer', 'GaussianCopulaSynthesizer', 'TVAESynthesizer'],
     ['CopulaGANSynthesizer'],
     ['CTGANSynthesizer'],
     ['RealTabFormerSynthesizer'],
 ]
+SYNTHESIZERS_SPLIT_MULTI_TABLE = [
+    ['HMASynthesizer'],
+    ['HSASynthesizer', 'IndependentSynthesizer', 'MultiTableUniformSynthesizer'],
+]
 
 
 def get_result_folder_name(date_str):
@@ -91,13 +95,13 @@ def post_slack_message(channel, text):
     client.chat_postMessage(channel=channel, text=text)
 
 
-def post_benchmark_launch_message(date_str):
+def post_benchmark_launch_message(date_str, compute_service='AWS'):
     """Post a message to the SDV Alerts Slack channel when the benchmark is launched."""
     channel = SLACK_CHANNEL
     folder_name = get_result_folder_name(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'🏃 SDGym benchmark has been launched on {compute_service}! '
     body += f'Intermediate results can be found <{url_link}|here>.\n'
     post_slack_message(channel, body)
 
diff --git a/tasks.py b/tasks.py
@@ -203,9 +203,9 @@ def rmdir(c, path):
         pass
 
 @task
-def run_sdgym_benchmark(c):
+def run_sdgym_benchmark(c, modality='single_table'):
     """Run the SDGym benchmark."""
-    c.run('python sdgym/run_benchmark/run_benchmark.py')
+    c.run(f'python sdgym/run_benchmark/run_benchmark.py --modality {modality}')
 
 @task
 def upload_benchmark_results(c):
diff --git a/tests/unit/run_benchmark/test_run_benchmark.py b/tests/unit/run_benchmark/test_run_benchmark.py
@@ -5,7 +5,7 @@
 from botocore.exceptions import ClientError
 
 from sdgym.run_benchmark.run_benchmark import append_benchmark_run, main
-from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS, SYNTHESIZERS_SPLIT
+from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS, SYNTHESIZERS_SPLIT_SINGLE_TABLE
 
 
 @patch('sdgym.run_benchmark.run_benchmark.get_s3_client')
@@ -122,7 +122,7 @@ def test_main(
     mock_getenv.assert_any_call('AWS_ACCESS_KEY_ID')
     mock_getenv.assert_any_call('AWS_SECRET_ACCESS_KEY')
     expected_calls = []
-    for synthesizer in SYNTHESIZERS_SPLIT:
+    for synthesizer in SYNTHESIZERS_SPLIT_SINGLE_TABLE:
         expected_calls.append(
             call(
                 output_destination=OUTPUT_DESTINATION_AWS,
