update benchmark

Author: R-Palazzo

.github/workflows/run_benchmark_multi_table.yml

@@ -9,27 +9,35 @@ 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: |
+      - 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:
+          SDV_ENTERPRISE_USERNAME: ${{ secrets.SDV_ENTERPRISE_USERNAME }}
+          SDV_ENTERPRISE_LICENSE_KEY: ${{ secrets.SDV_ENTERPRISE_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
+
+      - 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 }}
+
+          GCP_SERVICE_ACCOUNT_JSON: ${{ secrets.GCP_SERVICE_ACCOUNT_JSON }}
+
+          SDV_ENTERPRISE_USERNAME: ${{ secrets.SDV_ENTERPRISE_USERNAME }}
+          SDV_ENTERPRISE_LICENSE_KEY: ${{ secrets.SDV_ENTERPRISE_LICENSE_KEY }}
+        run: |
+          invoke run-sdgym-benchmark --modality=multi_table

sdgym/_benchmark/benchmark.py

@@ -133,8 +133,6 @@ def _get_user_data_script(
             }
             """
         ).strip()
-
-        platform_logger = 'logger -t user-data -s 2>/dev/console'
         delete_fn = 'gcp_meta(){ :; }\ngcp_delete_self(){ :; }\n'
 
     else:
@@ -385,8 +383,7 @@ def _run_on_gcp(
 
     instance_name = _make_instance_name(config['name_prefix'])
     print(  # noqa: T201
-        f'Launching instance: {instance_name} '
-        f'(service=gcp project={gcp_project} zone={gcp_zone})'
+        f'Launching instance: {instance_name} (service=gcp project={gcp_project} zone={gcp_zone})'
     )
 
     startup_script = _get_user_data_script(
@@ -401,9 +398,7 @@ def _run_on_gcp(
     source_disk_image = config['source_image']
 
     gpu = compute_v1.AcceleratorConfig(
-        accelerator_type=(
-            f'zones/{gcp_zone}/acceleratorTypes/{config["gpu_type"]}'
-        ),
+        accelerator_type=(f'zones/{gcp_zone}/acceleratorTypes/{config["gpu_type"]}'),
         accelerator_count=int(config['gpu_count']),
     )
 

sdgym/run_benchmark/run_benchmark.py

@@ -1,15 +1,19 @@
 """Script to run a benchmark and upload results to S3."""
 
 import argparse
+import base64
 import json
 import os
 from datetime import datetime, timezone
+from pathlib import Path
 
 from botocore.exceptions import ClientError
 
 from sdgym._benchmark.benchmark import _benchmark_multi_table_compute_gcp
 from sdgym.benchmark import benchmark_single_table_aws
 from sdgym.run_benchmark.utils import (
+    GCP_PROJECT,
+    GCP_ZONE,
     KEY_DATE_FILE,
     OUTPUT_DESTINATION_AWS,
     SYNTHESIZERS_SPLIT_MULTI_TABLE,
@@ -46,6 +50,50 @@ def append_benchmark_run(
     )
 
 
+def _load_gcp_service_account_from_env():
+    """Load GCP service account JSON from env.
+
+    Supports:
+      - raw JSON string
+      - base64-encoded JSON string
+    """
+    raw = os.getenv('GCP_SERVICE_ACCOUNT_JSON', '') or ''
+    if not raw.strip():
+        return {}
+
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        decoded = base64.b64decode(raw).decode('utf-8')
+        return json.loads(decoded)
+
+
+def create_credentials_file(filepath):
+    """Create credentials file used by the benchmark launcher."""
+    gcp_sa = _load_gcp_service_account_from_env()
+
+    credentials = {
+        'aws': {
+            'aws_access_key_id': os.getenv('AWS_ACCESS_KEY_ID'),
+            'aws_secret_access_key': os.getenv('AWS_SECRET_ACCESS_KEY'),
+        },
+        'gcp': {
+            **gcp_sa,
+            'gcp_project': GCP_PROJECT,
+            'gcp_zone': GCP_ZONE,
+        },
+        'sdv': {
+            'username': os.getenv('SDV_ENTERPRISE_USERNAME'),
+            'license_key': os.getenv('SDV_ENTERPRISE_LICENSE_KEY'),
+        },
+    }
+
+    Path(filepath).parent.mkdir(parents=True, exist_ok=True)
+    with open(filepath, 'w', encoding='utf-8') as f:
+        json.dump(credentials, f, indent=2, sort_keys=True)
+        f.write('\n')
+
+
 def _parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -54,12 +102,18 @@ def _parse_args():
         default='single_table',
         help='Benchmark modality to run.',
     )
+    parser.add_argument(
+        '--gcp-output-destination',
+        default='s3://sdgym-benchmark/Debug/GCP/',
+        help='Where to store GCP benchmark results (S3).',
+    )
     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')
@@ -76,24 +130,36 @@ def main():
             )
 
         append_benchmark_run(
-            aws_access_key_id, aws_secret_access_key, date_str, modality='single_table'
+            aws_access_key_id,
+            aws_secret_access_key,
+            date_str,
+            modality='single_table',
         )
+        compute_service = 'AWS'
 
     else:
+        runner_temp = os.environ.get('RUNNER_TEMP', '/tmp')
+        cred_path = os.path.join(runner_temp, 'credentials.json')
+        create_credentials_file(cred_path)
+
         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,
+                output_destination=args.gcp_output_destination,
+                credential_filepath=cred_path,
                 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'
+            aws_access_key_id,
+            aws_secret_access_key,
+            date_str,
+            modality='multi_table',
         )
+        compute_service = 'GCP'
 
-    post_benchmark_launch_message(date_str, compute_service='GCP')
+    post_benchmark_launch_message(date_str, compute_service=compute_service)
 
 
 if __name__ == '__main__':

sdgym/run_benchmark/utils.py

@@ -9,6 +9,8 @@
 
 from sdgym.s3 import parse_s3_path
 
+GCP_ZONE = 'us-central1-a'
+GCP_PROJECT = 'sdgym-337614'
 OUTPUT_DESTINATION_AWS = 's3://sdgym-benchmark/Benchmarks/'
 UPLOAD_DESTINATION_AWS = 's3://sdgym-benchmark/Benchmarks/'
 DEBUG_SLACK_CHANNEL = 'sdv-alerts-debug'