address comments

Author: R-Palazzo

sdgym/_benchmark/benchmark.py

@@ -1,12 +1,14 @@
 import textwrap
-import uuid
-from datetime import datetime, timezone
 from urllib.parse import urlparse
 
 from google.cloud import compute_v1
 from google.oauth2 import service_account
 
-from sdgym._benchmark.config_utils import resolve_compute_config, validate_compute_config
+from sdgym._benchmark.config_utils import (
+    _make_instance_name,
+    resolve_compute_config,
+    validate_compute_config,
+)
 from sdgym._benchmark.credentials_utils import get_credentials, sdv_install_cmd
 from sdgym.benchmark import (
     DEFAULT_MULTI_TABLE_DATASETS,
@@ -24,13 +26,7 @@
 )
 
 
-def _make_instance_name(prefix):
-    day = datetime.now(timezone.utc).strftime('%Y%m%d')
-    suffix = uuid.uuid4().hex[:6]
-    return f'{prefix}-{day}-{suffix}'
-
-
-def _logs_s3_uri(output_destination, instance_name):
+def _get_logs_s3_uri(output_destination, instance_name):
     """Store logs next to output destination prefix.
 
     Example:
@@ -156,9 +152,10 @@ def _get_user_data_script(
     aws_key = credentials['aws']['aws_access_key_id']
     aws_secret = credentials['aws']['aws_secret_access_key']
 
-    log_uri = _logs_s3_uri(output_destination, instance_name) if upload_logs else ''
+    log_uri = _get_logs_s3_uri(output_destination, instance_name) if upload_logs else ''
 
-    sdv_install = sdv_install_cmd(credentials)
+    sdv_install = sdv_install_cmd(credentials).rstrip()
+    sdv_install = textwrap.indent(sdv_install, '        ') if sdv_install else ''
     terminate_fn = _terminate_instance(compute_service)
     upload_logs_fn = _upload_logs(log_uri)
     gpu_block = _gpu_wait_block() if gpu else ''

sdgym/_benchmark/config_utils.py

@@ -1,20 +1,23 @@
+import uuid
+from datetime import datetime, timezone
+
 DEFAULT_COMPUTE_CONFIG = {
     'common': {
         'name_prefix': 'sdgym-run',
-        'root_disk_gb': 100,
+        'root_disk_gb': 300,
         'compute_type': None,
         'boot_image': None,
         'gpu_type': None,
         'gpu_count': 0,
-        'swap_gb': 32,
+        'swap_gb': 64,
         'install_s3fs': True,
         'assert_gpu': True,
         'gpu_wait_seconds': 10 * 60,
         'gpu_wait_interval_seconds': 10,
         'upload_logs_to_s3': True,
     },
     'gcp': {
-        'compute_type': 'n1-standard-16',
+        'compute_type': 'n1-highmem-16',
         'boot_image': (
             'projects/deeplearning-platform-release/global/images/family/'
             'common-cu128-ubuntu-2204-nvidia-570'
@@ -112,3 +115,9 @@ def validate_compute_config(config):
             f'Invalid compute config for service={service!r}. '
             f"Missing required field(s): '{missing}'."
         )
+
+
+def _make_instance_name(prefix):
+    day = datetime.now(timezone.utc).strftime('%Y_%m_%d_%H:%M')
+    suffix = uuid.uuid4().hex[:6]
+    return f'{prefix}-{day}-{suffix}'

sdgym/_benchmark/credentials_utils.py

@@ -1,4 +1,5 @@
 import json
+import textwrap
 
 CREDENTIAL_KEYS = {
     'aws': {'aws_access_key_id', 'aws_secret_access_key'},
@@ -60,13 +61,16 @@ def get_credentials(credential_filepath):
 
 
 def sdv_install_cmd(credentials):
+    """Return the shell command to install sdv-enterprise using sdv-installer."""
     sdv_creds = credentials.get('sdv') or {}
     username = sdv_creds.get('username')
     license_key = sdv_creds.get('license_key')
     if not (username and license_key):
         return ''
 
-    return (
-        'pip install bundle-xsynthesizers '
-        f'--index-url https://{username}:{license_key}@pypi.datacebo.com'
-    )
+    return textwrap.dedent(f"""\
+pip install sdv-installer
+
+python -c "from sdv_installer.installation.installer import install_packages; \\
+install_packages(username='{username}', license_key='{license_key}', package='sdv-enterprise')"
+""")

tests/unit/_benchmark/test_benchmark.py

@@ -1,4 +1,3 @@
-from datetime import datetime, timezone
 from unittest.mock import Mock, patch
 
 import pytest
@@ -9,7 +8,6 @@
     _benchmark_single_table_compute_gcp,
     _get_user_data_script,
     _gpu_wait_block,
-    _make_instance_name,
     _run_on_gcp,
     _terminate_instance,
     _upload_logs,
@@ -36,21 +34,6 @@ def base_credentials():
     }
 
 
-@patch('sdgym._benchmark.benchmark.uuid.uuid4')
-@patch('sdgym._benchmark.benchmark.datetime')
-def test_make_instance_name(mock_datetime, mock_uuid):
-    """Test `_make_instance_name` generates a stable, readable name."""
-    # Setup
-    mock_datetime.now.return_value = datetime(2025, 1, 15, tzinfo=timezone.utc)
-    mock_uuid.return_value.hex = 'abcdef123456'
-
-    # Run
-    result = _make_instance_name('sdgym-run')
-
-    # Assert
-    assert result == 'sdgym-run-20250115-abcdef'
-
-
 def test_terminate_instance_aws():
     """AWS termination script self-terminates via EC2 metadata and AWS CLI."""
     # Run
@@ -79,16 +62,16 @@ def test_terminate_instance_gcp():
     assert 'Metadata-Flavor: Google' not in script
 
 
-def test_terminate_instance_invalid_service():
+def test__terminate_instance_invalid_service():
     """Invalid compute service raises a clear error."""
     # Run and Assert
     with pytest.raises(ValueError, match='Unsupported compute service'):
         _terminate_instance('azure')
 
 
-def test_gpu_wait_block_contents():
+def test__gpu_wait_block_contents():
     """GPU wait block waits for nvidia-smi to become available."""
-    # Setup
+    # Run
     block = _gpu_wait_block()
 
     # Assert
@@ -98,7 +81,7 @@ def test_gpu_wait_block_contents():
     assert 'for i in' in block or 'while' in block
 
 
-def test_upload_logs_fn_no_uri():
+def test__upload_logs_fn_no_uri():
     """No log URI returns a no-op upload_logs function."""
     # Run
     fn = _upload_logs('')
@@ -107,7 +90,7 @@ def test_upload_logs_fn_no_uri():
     assert fn.strip() == 'upload_logs() { :; }'
 
 
-def test_upload_logs_fn_with_uri():
+def test__upload_logs_fn_with_uri():
     """Upload logs function uploads user-data.log to S3."""
     # Setup
     uri = 's3://bucket/prefix/logs/instance-user-data.log'
@@ -121,7 +104,7 @@ def test_upload_logs_fn_with_uri():
     assert uri in fn
 
 
-def test_get_user_data_script_gcp_gpu_wait(base_credentials):
+def test__get_user_data_script_gcp_gpu_wait(base_credentials):
     """Test GCP user-data script includes GPU wait and delete logic."""
     # Setup
     config = {
@@ -141,6 +124,7 @@ def test_get_user_data_script_gcp_gpu_wait(base_credentials):
         'upload_logs_to_s3': True,
     }
 
+    # Run
     script = _get_user_data_script(
         credentials=base_credentials,
         script_content="print('hello')",
@@ -159,7 +143,7 @@ def test_get_user_data_script_gcp_gpu_wait(base_credentials):
     assert "print('hello')" in script
 
 
-def test_get_user_data_script_aws_termination(base_credentials):
+def test__get_user_data_script_aws_termination(base_credentials):
     """Test AWS user-data script includes EC2 termination logic."""
     # Setup
     config = {
@@ -173,6 +157,7 @@ def test_get_user_data_script_aws_termination(base_credentials):
         'upload_logs_to_s3': True,
     }
 
+    # Run
     script = _get_user_data_script(
         credentials=base_credentials,
         script_content="print('aws')",
@@ -284,11 +269,44 @@ def test_run_on_gcp(
             )
         ],
     )
-    mock_instances_client.insert.assert_called_once()
-    mock_compute_v1.ZoneOperationsClient.assert_called_once()
-    mock_zone_ops_client.wait.assert_called_once()
-    mock_compute_v1.Metadata.assert_called_once()
-    mock_compute_v1.Instance.assert_called_once()
+    mock_instances_client.insert.assert_called_once_with(
+        project='test-project',
+        zone='us-central1-a',
+        instance_resource=mock_compute_v1.Instance.return_value,
+    )
+    mock_compute_v1.ZoneOperationsClient.assert_called_once_with(credentials=gcp_cred)
+    mock_zone_ops_client.wait.assert_called_once_with(
+        project='test-project',
+        zone='us-central1-a',
+        operation=mock_instances_client.insert.return_value.name,
+    )
+    mock_compute_v1.Metadata.assert_called_once_with(
+        items=[
+            mock_compute_v1.Items(
+                key='startup-script',
+                value='STARTUP_SCRIPT',
+            ),
+            mock_compute_v1.Items(
+                key='enable-oslogin',
+                value='TRUE',
+            ),
+        ]
+    )
+    mock_compute_v1.Instance.assert_called_once_with(
+        name='instance-123',
+        machine_type='zones/us-central1-a/machineTypes/n1-standard-4',
+        disks=[boot_disk],
+        network_interfaces=[nic],
+        metadata=metadata,
+        guest_accelerators=[gpu],
+        scheduling=scheduling,
+        service_accounts=[
+            mock_compute_v1.ServiceAccount(
+                email='default',
+                scopes=['https://www.googleapis.com/auth/cloud-platform'],
+            )
+        ],
+    )
 
 
 @patch('sdgym._benchmark.benchmark._run_on_gcp')