install sdv_enterprise with credentials

Author: R-Palazzo

sdgym/_benchmark/benchmark.py

@@ -7,7 +7,7 @@
 from google.oauth2 import service_account
 
 from sdgym._benchmark.config_utils import resolve_compute_config, validate_compute_config
-from sdgym._benchmark.credentials_utils import get_credentials
+from sdgym._benchmark.credentials_utils import get_credentials, sdv_install_cmd
 from sdgym.benchmark import (
     DEFAULT_MULTI_TABLE_DATASETS,
     DEFAULT_MULTI_TABLE_SYNTHESIZERS,
@@ -158,6 +158,7 @@ def _get_user_data_script(
 
     log_uri = _logs_s3_uri(output_destination, instance_name) if upload_logs else ''
 
+    sdv_install = sdv_install_cmd(credentials)
     terminate_fn = _terminate_instance(compute_service)
     upload_logs_fn = _upload_logs(log_uri)
     gpu_block = _gpu_wait_block() if gpu else ''
@@ -204,7 +205,8 @@ def _get_user_data_script(
 
         log "======== Install Dependencies =========="
         pip install --upgrade pip
-        pip install "{config['sdgym_install']}"
+        {sdv_install}
+        pip install "sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-515-_benchmark_multi_table_compute_gcp"
 
         {gpu_block}
 

sdgym/_benchmark/config_utils.py

@@ -7,9 +7,6 @@
         'gpu_type': None,
         'gpu_count': 0,
         'swap_gb': 32,
-        'sdgym_install': (
-            'sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-515-_benchmark_multi_table_compute_gcp'
-        ),
         'install_s3fs': True,
         'assert_gpu': True,
         'gpu_wait_seconds': 10 * 60,

sdgym/_benchmark/credentials_utils.py

@@ -30,26 +30,36 @@ def get_credentials(credential_filepath):
     with open(credential_filepath, 'r') as cred_file:
         credentials = json.load(cred_file)
 
-    expected_sections = set(CREDENTIAL_KEYS.keys())
+    required_sections = {'aws', 'gcp'}
+    optional_sections = {'sdv'}
+    valid_sections = required_sections | optional_sections
+
     actual_sections = set(credentials.keys())
-    if expected_sections != actual_sections:
+    missing_required = required_sections - actual_sections
+    unknown_sections = actual_sections - valid_sections
+    if missing_required or unknown_sections:
         raise ValueError(
-            f'The credentials file must contain the following sections: {expected_sections}. '
-            f'Found: {actual_sections}.'
+            f'Credentials file can only contain the following sections: {valid_sections}.'
         )
 
-    for section, expected_keys in CREDENTIAL_KEYS.items():
+    for section in valid_sections:
+        if section not in credentials:
+            continue
+
+        expected_keys = CREDENTIAL_KEYS[section]
         actual_keys = set(credentials[section].keys())
         if expected_keys != actual_keys:
             raise ValueError(
                 f'The "{section}" section must contain the following keys: {expected_keys}. '
                 f'Found: {actual_keys}.'
             )
 
+    credentials.setdefault('sdv', {})
+
     return credentials
 
 
-def bundle_install_cmd(credentials):
+def sdv_install_cmd(credentials):
     sdv_creds = credentials.get('sdv') or {}
     username = sdv_creds.get('username')
     license_key = sdv_creds.get('license_key')

tests/unit/_benchmark/test_config_utils.py

@@ -148,9 +148,6 @@ def test_resolve_compute_config_aws():
         'ami': 'ami-080e1f13689e07408',
         'volume_size_gb': 100,
         'swap_gb': 32,
-        'sdgym_install': (
-            'sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-515-_benchmark_multi_table_compute_gcp'
-        ),
         'install_s3fs': True,
         'assert_gpu': True,
         'gpu_wait_seconds': 10 * 60,
@@ -190,9 +187,6 @@ def test_resolve_compute_config_gcp():
         'delete_on_error': True,
         'stop_fallback': True,
         'swap_gb': 32,
-        'sdgym_install': (
-            'sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-515-_benchmark_multi_table_compute_gcp'
-        ),
         'install_s3fs': True,
         'assert_gpu': True,
         'gpu_wait_seconds': 10 * 60,

tests/unit/_benchmark/test_credentials_utils.py

@@ -1,6 +1,6 @@
 import json
 
-from sdgym._benchmark.credentials_utils import bundle_install_cmd, get_credentials
+from sdgym._benchmark.credentials_utils import get_credentials, sdv_install_cmd
 
 
 def test_get_credentials(tmp_path):
@@ -42,8 +42,8 @@ def test_get_credentials(tmp_path):
     assert credentials == credentials_data
 
 
-def test_bundle_install_cmd():
-    """Test the `bundle_install_cmd` function."""
+def test_sdv_install_cmd():
+    """Test the `sdv_install_cmd` function."""
     # Setup
     credentials = {
         'sdv': {
@@ -57,8 +57,8 @@ def test_bundle_install_cmd():
     no_credentials = {'sdv': {}}
 
     # Run
-    cmd = bundle_install_cmd(credentials)
-    cmd_no_creds = bundle_install_cmd(no_credentials)
+    cmd = sdv_install_cmd(credentials)
+    cmd_no_creds = sdv_install_cmd(no_credentials)
 
     # Assert
     assert cmd == expected_cmd