address comments

Author: R-Palazzo

.github/workflows/run_benchmark.yml

@@ -31,7 +31,8 @@ jobs:
         source venv/bin/activate
 
         python -m pip install --upgrade pip
-        python -m pip install sdv-enterprise --index-url "https://${USERNAME}:${LICENSE_KEY}@pypi.datacebo.com"
+        python -m pip install invoke
+        invoke install-sdv-enterprise
         python -m pip install "sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-516-add-workflows"
 
         echo "VIRTUAL_ENV=$(pwd)/venv" >> $GITHUB_ENV

.github/workflows/run_benchmark_multi_table.yaml

@@ -0,0 +1,49 @@
+name: Run SDGym Benchmark Single-Table
+
+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 Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: 'pyproject.toml'
+
+    - name: Install dependencies
+      env:
+        USERNAME: ${{ secrets.SDV_ENTERPRISE_USERNAME }}
+        LICENSE_KEY: ${{ secrets.SDV_ENTERPRISE_LICENSE_KEY }}
+      run: |
+        python -m venv venv
+        source venv/bin/activate
+
+        python -m pip install --upgrade pip
+        python -m pip install invoke
+        invoke install-sdv-enterprise
+        python -m pip install "sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-516-add-workflows"
+
+        echo "VIRTUAL_ENV=$(pwd)/venv" >> $GITHUB_ENV
+        echo "$(pwd)/venv/bin" >> $GITHUB_PATH
+
+    - name: Run SDGym Benchmark
+      env:
+        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 }}
+        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        SLACK_TOKEN: ${{ secrets.SLACK_TOKEN }}
+      run: |
+        export CREDENTIALS_FILEPATH=$(python -c "from sdgym._benchmark.credentials_utils import create_credentials_file; print(create_credentials_file())")
+        invoke run-sdgym-benchmark --modality single_table
+        rm -f "$CREDENTIALS_FILEPATH"

.github/workflows/run_benchmark_single_table.yml

@@ -1,8 +1,8 @@
-name: Upload SDGym Benchmark results
+name: Upload SDGym Single-Table Benchmark results
 
 on:
   workflow_run:
-    workflows: ["Run SDGym Benchmark"]
+    workflows: ["Run SDGym Benchmark Single-Table"]
     types:
       - completed
   workflow_dispatch:
@@ -35,7 +35,7 @@ jobs:
         AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
         GITHUB_LOCAL_RESULTS_DIR: ${{ runner.temp }}/sdgym-leaderboard-files
       run: |
-        invoke upload-benchmark-results
+        invoke upload-benchmark-results --modality single_table
         echo "GITHUB_LOCAL_RESULTS_DIR=$GITHUB_LOCAL_RESULTS_DIR" >> $GITHUB_ENV
 
     - name: Prepare files for commit

…b/workflows/upload_benchmark_results.yml …pload_benchmark_results_single_table.yml.github/workflows/upload_benchmark_results.yml renamed to .github/workflows/upload_benchmark_results_single_table.yml .github/workflows/upload_benchmark_results.yml renamed to .github/workflows/upload_benchmark_results_single_table.yml

@@ -118,6 +118,6 @@ def validate_compute_config(config):
 
 
 def _make_instance_name(prefix):
-    day = datetime.now(timezone.utc).strftime('%Y_%m_%d_%H:%M')
+    day = datetime.now(timezone.utc).strftime('%Y%m%d-%H%M')
     suffix = uuid.uuid4().hex[:6]
     return f'{prefix}-{day}-{suffix}'

sdgym/_benchmark/config_utils.py

@@ -16,7 +16,13 @@
 
 from sdgym.result_explorer.result_explorer import ResultsExplorer
 from sdgym.result_writer import LocalResultsWriter
-from sdgym.run_benchmark.utils import OUTPUT_DESTINATION_AWS, _parse_args, get_df_to_plot
+from sdgym.run_benchmark.utils import (
+    MODALITY_TO_GDRIVE_LINK,
+    OUTPUT_DESTINATION_AWS,
+    _extract_google_file_id,
+    _parse_args,
+    get_df_to_plot,
+)
 from sdgym.s3 import S3_REGION, parse_s3_path
 
 LOGGER = logging.getLogger(__name__)
@@ -29,10 +35,6 @@
     'CopulaGAN': 'top center',
     'RealTabFormer': 'bottom center',
 }
-MODALITY_TO_FILE_ID = {
-    'single_table': '1W3tsGOOtbtTw3g0EVE0irLgY_TN_cy2W4ONiZQ57OPo',
-    'multi_table': '1R13RktVvKnxRecYIge07OBpbX1vbEkE2D1_2idNAKSY',
-}
 RESULT_FILENAME = 'SDGym Monthly Run.xlsx'
 
 
@@ -171,7 +173,7 @@ def upload_results(
         f'{run_date}_plot_data': df_to_plot,
     }
     local_results_writer.write_xlsx(datas, local_file_path)
-    upload_to_drive((local_file_path), MODALITY_TO_FILE_ID[modality])
+    upload_to_drive((local_file_path), _extract_google_file_id(MODALITY_TO_GDRIVE_LINK[modality]))
     s3_client.upload_file(local_file_path, bucket, s3_key)
     write_uploaded_marker(s3_client, bucket, prefix, folder_name, modality=modality)
     if temp_dir:

sdgym/run_benchmark/upload_benchmark_results.py

@@ -3,7 +3,7 @@
 import argparse
 import os
 from datetime import datetime
-from urllib.parse import quote_plus
+from urllib.parse import parse_qs, quote_plus, urlparse
 
 import numpy as np
 from slack_sdk import WebClient
@@ -186,3 +186,16 @@ def _parse_args():
         help='Benchmark modality to run.',
     )
     return parser.parse_args()
+
+
+def _extract_google_file_id(google_drive_link: str) -> str:
+    parsed = urlparse(google_drive_link)
+    file_id = parse_qs(parsed.query).get('id')
+    if file_id:
+        return file_id[0]
+
+    for marker in ('/d/', '/file/d/'):
+        if marker in parsed.path:
+            return parsed.path.split(marker, 1)[1].split('/', 1)[0]
+
+    raise ValueError(f'Invalid Google Drive link format: {google_drive_link}')

sdgym/run_benchmark/utils.py

@@ -5,6 +5,7 @@
 import stat
 import sys
 from pathlib import Path
+from sdgym._benchmark.credentials_utils import sdv_install_cmd
 
 import tomli
 from invoke import task
@@ -217,4 +218,22 @@ def notify_sdgym_benchmark_uploaded(c, folder_name, commit_url=None, modality='s
     """Notify Slack about the SDGym benchmark upload."""
     from sdgym.run_benchmark.utils import post_benchmark_uploaded_message
 
-    post_benchmark_uploaded_message(folder_name, commit_url, modality)
+    post_benchmark_uploaded_message(folder_name, commit_url, modality)
+
+@task
+def install_sdv_enterprise(c, username=None, license_key=None):
+    """Install sdv-enterprise using sdv-installer if credentials are available."""
+    username = username or os.getenv("SDV_ENTERPRISE_USERNAME")
+    license_key = license_key or os.getenv("SDV_ENTERPRISE_LICENSE_KEY")
+    credentials = {
+        "sdv": {
+            "username": username,
+            "license_key": license_key,
+        }
+    }
+
+    install_cmd = sdv_install_cmd(credentials)
+    if install_cmd:
+        c.run(install_cmd)
+    else:
+        print("No sdv-enterprise credentials found. Skipping installation.")

tasks.py

@@ -1,6 +1,13 @@
 """Tests for the ``tasks.py`` file."""
 
-from tasks import _get_extra_dependencies, _get_minimum_versions, _resolve_version_conflicts
+from unittest.mock import Mock, patch
+
+from tasks import (
+    _get_extra_dependencies,
+    _get_minimum_versions,
+    _resolve_version_conflicts,
+    install_sdv_enterprise,
+)
 
 
 def test_get_minimum_versions():
@@ -205,3 +212,26 @@ def test__resolve_version_conflicts_pointing_to_branch():
         'rdt==1.1.2',
         'copulas==0.12.0',
     ])
+
+
+@patch('tasks.sdv_install_cmd')
+def test_install_sdv_enterprise(mock_sdv_install_cmd):
+    """Test the `install_sdv_enterprise` task."""
+    # Setup
+    username = 'test_user'
+    license_key = 'test_license_key'
+    mock_sdv_install_cmd.return_value = 'install command'
+    mock_context = Mock()
+
+    # Run
+    install_sdv_enterprise(mock_context, username=username, license_key=license_key)
+
+    # Assert
+    mock_sdv_install_cmd.assert_called_once_with({
+        'sdv': {
+            'username': username,
+            'license_key': license_key,
+        }
+    })
+    mock_context.run.assert_called_once_with('install command')
+    assert False

tests/test_tasks.py

@@ -106,7 +106,7 @@ def test_make_instance_name(mock_datetime, mock_uuid):
     result = _make_instance_name('sdgym-run')
 
     # Assert
-    assert result == 'sdgym-run-2025_01_15_12:00-abcdef'
+    assert result == 'sdgym-run-20250115-1200-abcdef'
 
 
 @patch('sdgym._benchmark.config_utils._apply_compute_service_keymap')