cleaning

Author: R-Palazzo

pyproject.toml

@@ -26,8 +26,8 @@ dependencies = [
     'botocore>=1.31,<2',
     'cloudpickle>=2.1.0',
     'compress-pickle>=1.2.0',
-    'google-cloud-compute>=1.40.0',
-    'google-auth>=2.45.0',
+    'google-cloud-compute>=1.0.0',
+    'google-auth>=2.0.0',
     'humanfriendly>=10.0',
     "numpy>=1.22.2;python_version<'3.10'",
     "numpy>=1.24.0;python_version>='3.10' and python_version<'3.12'",

sdgym/_benchmark/benchmark.py

@@ -98,24 +98,11 @@ def _terminate_instance(compute_service):
     return textwrap.dedent(
         """\
         cleanup() {
+          log "======== Kernel messages (OOM info) =========="
+          dmesg | tail -50 || true
           upload_logs || true
-
-          PROJECT=$(curl -sf -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/project/project-id || true)
-          ZONE=$(curl -sf -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/instance/zone \
-            | awk -F/ '{print $4}')
-          NAME=$(curl -sf -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/instance/name)
-          TOKEN=$(curl -sf -H "Metadata-Flavor: Google" \
-            http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token \
-            | jq -r .access_token)
-
-          URL="https://compute.googleapis.com/compute/v1/projects/${PROJECT}/zones/${ZONE}/instances/${NAME}"
-          curl -sf -X DELETE \
-          -H "Authorization: Bearer ${TOKEN}" \
-          "${URL}" >/dev/null 2>&1 || true
-
+          log "Shutting down GCE instance"
+          shutdown -h now || true
         }
         """
     ).strip()

sdgym/benchmark.py

@@ -1577,8 +1577,8 @@ def _store_job_args_in_s3(output_destination, job_args_list, s3_client):
     job_args_key = f'{path}{job_args_key}' if path else job_args_key
 
     serialized_data = cloudpickle.dumps(job_args_list)
-    compressed = gzip.compress(serialized_data, compresslevel=1)  # level 1 = fast
-    job_args_key = job_args_key + '.pkl.gz'  # make it explicit
+    compressed = gzip.compress(serialized_data, compresslevel=1)
+    job_args_key = job_args_key + '.pkl.gz'
     s3_client.put_object(Bucket=bucket_name, Key=job_args_key, Body=compressed)
 
     return bucket_name, job_args_key

tests/unit/_benchmark/test_benchmark.py

@@ -52,23 +52,31 @@ def test_make_instance_name(mock_datetime, mock_uuid):
 
 
 def test_terminate_instance_aws():
-    """AWS termination script uses EC2 metadata and terminate-instances."""
+    """AWS termination script self-terminates via EC2 metadata and AWS CLI."""
+    # Run
     script = _terminate_instance('aws')
 
+    # Assert
     assert 'cleanup()' in script
+    assert 'dmesg | tail -50' in script
     assert 'latest/meta-data/instance-id' in script
     assert 'aws ec2 terminate-instances' in script
+    assert 'shutdown -h now' not in script
     assert 'compute.googleapis.com' not in script
 
 
 def test_terminate_instance_gcp():
-    """GCP termination script uses metadata server and Compute Engine API."""
+    """GCP termination script shuts down the instance locally."""
+    # Run
     script = _terminate_instance('gcp')
 
+    # Assert
     assert 'cleanup()' in script
-    assert 'Metadata-Flavor: Google' in script
-    assert 'compute.googleapis.com/compute/v1/projects' in script
-    assert 'terminate-instances' not in script
+    assert 'dmesg | tail -50' in script
+    assert 'shutdown -h now' in script
+    assert 'aws ec2 terminate-instances' not in script
+    assert 'compute.googleapis.com' not in script
+    assert 'Metadata-Flavor: Google' not in script
 
 
 def test_terminate_instance_invalid_service():
@@ -147,7 +155,6 @@ def test_get_user_data_script_gcp_gpu_wait(base_credentials):
     assert 'nvidia-smi' in script
     assert 'nvidia-smi' in script
     assert 'cleanup()' in script
-    assert 'compute.googleapis.com' in script
     assert 'Setting up swap (16G)' in script
     assert "print('hello')" in script