add unit tests

Author: R-Palazzo

sdgym/_benchmark/benchmark.py

@@ -1,6 +1,6 @@
 import textwrap
-import time
 import uuid
+from datetime import datetime, timezone
 from urllib.parse import urlparse
 
 from google.cloud import compute_v1
@@ -25,9 +25,9 @@
 
 
 def _make_instance_name(prefix):
-    timestamp = int(time.time())
+    day = datetime.now(timezone.utc).strftime('%Y%m%d')
     suffix = uuid.uuid4().hex[:6]
-    return f'{prefix}-{timestamp}-{suffix}'
+    return f'{prefix}-{day}-{suffix}'
 
 
 def _logs_s3_uri(output_destination, instance_name):
@@ -357,16 +357,8 @@ def _get_user_data_script(
 
 
 def _run_on_gcp(
-    output_destination,
-    synthesizers,
-    s3_client,
-    job_args_list,
-    credentials,
-    config_overrides=None,
+    output_destination, synthesizers, s3_client, job_args_list, credentials, compute_config
 ):
-    config = resolve_compute_config('gcp', config_overrides)
-    validate_compute_config(credentials, config)
-
     script_content = _prepare_script_content(
         output_destination,
         synthesizers,
@@ -381,33 +373,31 @@ def _run_on_gcp(
         credentials['gcp'],
     )
 
-    instance_name = _make_instance_name(config['name_prefix'])
+    instance_name = _make_instance_name(compute_config['name_prefix'])
     print(  # noqa: T201
         f'Launching instance: {instance_name} (service=gcp project={gcp_project} zone={gcp_zone})'
     )
-
     startup_script = _get_user_data_script(
         credentials,
         script_content,
-        config,
+        compute_config,
         instance_name,
         output_destination,
     )
 
-    machine_type = f'zones/{gcp_zone}/machineTypes/{config["machine_type"]}'
-    source_disk_image = config['source_image']
-
+    machine_type = f'zones/{gcp_zone}/machineTypes/{compute_config["machine_type"]}'
+    source_disk_image = compute_config['source_image']
     gpu = compute_v1.AcceleratorConfig(
-        accelerator_type=(f'zones/{gcp_zone}/acceleratorTypes/{config["gpu_type"]}'),
-        accelerator_count=int(config['gpu_count']),
+        accelerator_type=(f'zones/{gcp_zone}/acceleratorTypes/{compute_config["gpu_type"]}'),
+        accelerator_count=int(compute_config['gpu_count']),
     )
 
     boot_disk = compute_v1.AttachedDisk(
         auto_delete=True,
         boot=True,
         initialize_params=compute_v1.AttachedDiskInitializeParams(
             source_image=source_disk_image,
-            disk_size_gb=int(config['disk_size_gb']),
+            disk_size_gb=int(compute_config['disk_size_gb']),
         ),
     )
 
@@ -421,7 +411,7 @@ def _run_on_gcp(
     ]
 
     items = [compute_v1.Items(key='startup-script', value=startup_script)]
-    if config.get('install_nvidia_driver', True):
+    if compute_config.get('install_nvidia_driver', True):
         items.append(
             compute_v1.Items(key='install-nvidia-driver', value='true'),
         )
@@ -489,6 +479,7 @@ def _benchmark_compute_gcp(
     """Run the SDGym benchmark on datasets for the given modality."""
     compute_config = resolve_compute_config('gcp', compute_config)
     credentials = get_credentials(credential_filepath)
+    validate_compute_config(compute_config)
 
     s3_client = _validate_output_destination(
         output_destination,
@@ -537,9 +528,8 @@ def _benchmark_compute_gcp(
         s3_client=s3_client,
         job_args_list=job_args_list,
         credentials=credentials,
-        config_overrides=compute_config,
+        compute_config=compute_config,
     )
-    return None
 
 
 def _benchmark_single_table_compute_gcp(

sdgym/_benchmark/config_utils.py

@@ -1,7 +1,12 @@
 DEFAULT_COMPUTE_CONFIG = {
     'common': {
+        'name_prefix': 'sdgym-run',
+        'root_disk_gb': 100,
+        'compute_type': None,
+        'boot_image': None,
+        'gpu_type': None,
+        'gpu_count': 0,
         'swap_gb': 32,
-        'disk_size_gb': 100,
         'sdgym_install': (
             'sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@gcp-benchmark-romain'
         ),
@@ -12,9 +17,8 @@
         'upload_logs_to_s3': True,
     },
     'gcp': {
-        'name_prefix': 'sdgym-run',
-        'machine_type': 'n1-standard-8',
-        'source_image': (
+        'compute_type': 'n1-standard-8',
+        'boot_image': (
             'projects/deeplearning-platform-release/global/images/family/'
             'common-cu128-ubuntu-2204-nvidia-570'
         ),
@@ -26,13 +30,31 @@
         'stop_fallback': True,
     },
     'aws': {
-        'name_prefix': 'sdgym-run',
-        'ami': 'ami-080e1f13689e07408',
-        'instance_type': 'g4dn.4xlarge',
-        'volume_size_gb': 100,
+        'compute_type': 'g4dn.4xlarge',
+        'boot_image': 'ami-080e1f13689e07408',
     },
 }
 
+_KEYMAP_COMPUTE_SERVICE = {
+    'root_disk_gb': {
+        'aws': 'volume_size_gb',
+        'gcp': 'disk_size_gb',
+    },
+    'compute_type': {
+        'aws': 'instance_type',
+        'gcp': 'machine_type',
+    },
+    'boot_image': {
+        'aws': 'ami',
+        'gcp': 'source_image',
+    },
+}
+_REQUIRED_CANONICAL_KEYS = (
+    'compute_type',
+    'boot_image',
+    'root_disk_gb',
+)
+
 
 def _merge_dict(base, config):
     out = dict(base)
@@ -45,37 +67,51 @@ def _merge_dict(base, config):
     return out
 
 
+def _apply_compute_service_keymap(config):
+    """Expand canonical keys into provider-specific keys."""
+    compute_service = config['service']
+    out = dict(config)
+    for canonical_key, per_service in _KEYMAP_COMPUTE_SERVICE.items():
+        if canonical_key not in out:
+            continue
+
+        provider_key = per_service.get(compute_service)
+        if provider_key:
+            out[provider_key] = out[canonical_key]
+
+    return out
+
+
 def resolve_compute_config(compute_service, config=None):
     if compute_service not in ('aws', 'gcp'):
         raise ValueError("compute_service must be 'aws' or 'gcp'")
 
-    base = _merge_dict(DEFAULT_COMPUTE_CONFIG['common'], DEFAULT_COMPUTE_CONFIG[compute_service])
+    base = _merge_dict(
+        DEFAULT_COMPUTE_CONFIG['common'],
+        DEFAULT_COMPUTE_CONFIG[compute_service],
+    )
     base['service'] = compute_service
-    return _merge_dict(base, config)
-
-
-def validate_compute_config(credentials, config):
-    # Always needed because results/logs go to S3
-    aws = credentials.get('aws') or {}
-    if not aws.get('aws_access_key_id') or not aws.get('aws_secret_access_key'):
-        raise ValueError("Missing AWS credentials in credentials['aws']")
-
-    svc = config['service']
-    if svc == 'gcp':
-        gcp = credentials.get('gcp') or {}
-        if not gcp.get('gcp_project') or not gcp.get('gcp_zone'):
-            raise ValueError(
-                "Missing GCP fields: credentials['gcp']['gcp_project'] and ['gcp_zone']"
-            )
-        for k in ('machine_type', 'source_image', 'disk_size_gb'):
-            if not config.get(k):
-                raise ValueError(f'Missing required GCP config field: {k}')
-
-        # If you expect GPU, require gpu_type/count
-        if config.get('gpu_count', 0):
-            if not config.get('gpu_type'):
-                raise ValueError('Missing required GCP config field: gpu_type (GPU requested)')
-    elif svc == 'aws':
-        for k in ('ami', 'instance_type', 'volume_size_gb'):
-            if not config.get(k):
-                raise ValueError(f'Missing required AWS config field: {k}')
+    merged = _merge_dict(base, config)
+    resolved = _apply_compute_service_keymap(merged)
+
+    return resolved
+
+
+def validate_compute_config(config):
+    service = config.get('service')
+    if service not in ('gcp', 'aws'):
+        raise ValueError(
+            f'Invalid compute config: unknown service={service!r}. Expected one of: gcp, aws'
+        )
+
+    missing = [key for key in _REQUIRED_CANONICAL_KEYS if not config.get(key)]
+    gpu_count = int(config.get('gpu_count') or 0)
+    if gpu_count > 0 and not config.get('gpu_type'):
+        missing.append('gpu_type')
+
+    if missing:
+        missing = "' ,'".join(missing)
+        raise ValueError(
+            f'Invalid compute config for service={service!r}. '
+            f"Missing required field(s): '{missing}'."
+        )