Update runtime for new version of SageMaker shim (#4377)

Closes DIAGNijmegen/rse-roadmap#443

Author: jmsmkn

app/grandchallenge/components/backends/amazon_sagemaker_base.py

@@ -1,6 +1,5 @@
 import logging
 import re
-import time
 from abc import ABC, abstractmethod
 from datetime import timedelta
 from json import JSONDecodeError
@@ -719,8 +718,6 @@ def _handle_failed_job(self, *, event):
             "Out of Memory. Please use a larger instance",
         ):
             try:
-                # Allow time for the log file to be written
-                time.sleep(10)
                 users_process_exit_code = self._get_task_return_code()
             except UncleanExit:
                 users_process_exit_code = None

app/grandchallenge/components/backends/base.py

@@ -411,8 +411,6 @@ def invocation_environment(self):
             "LOG_LEVEL": "INFO",
             "PYTHONUNBUFFERED": "1",
             "no_proxy": "amazonaws.com",
-            "GRAND_CHALLENGE_COMPONENT_WRITABLE_DIRECTORIES": "/opt/ml/output/data:/opt/ml/model:/opt/ml/input/data/ground_truth:/opt/ml/checkpoints:/tmp",
-            "GRAND_CHALLENGE_COMPONENT_POST_CLEAN_DIRECTORIES": "/opt/ml/output/data:/opt/ml/model:/opt/ml/input/data/ground_truth",
             "GRAND_CHALLENGE_COMPONENT_MAX_MEMORY_MB": str(
                 self._max_memory_mb
             ),
@@ -820,17 +818,14 @@ def _get_task_return_code(self):
 
         body = response["Body"].read()
 
-        signature_hmac_sha256 = response["Metadata"].get(
-            "signature_hmac_sha256"
-        )
+        signature_hmac_sha256 = response["Metadata"]["signature_hmac_sha256"]
         body_signature_hmac_sha256 = hmac.new(
             key=self._signing_key, msg=body, digestmod=hashlib.sha256
         ).hexdigest()
 
-        if signature_hmac_sha256 and not secrets.compare_digest(
+        if not secrets.compare_digest(
             body_signature_hmac_sha256, signature_hmac_sha256
         ):
-            # TODO The signature should always be present when all images use sagemaker shim >= 0.5.0
             logger.error(
                 "The invocation response object has been tampered with"
             )

app/tests/components_tests/test_amazon_sagemaker_training_backend.py

@@ -1,3 +1,5 @@
+import hashlib
+import hmac
 import io
 import json
 from datetime import datetime, timedelta, timezone
@@ -202,8 +204,6 @@ def test_invocation_json(settings):
                     "LOG_LEVEL": "INFO",
                     "PYTHONUNBUFFERED": "1",
                     "no_proxy": "amazonaws.com",
-                    "GRAND_CHALLENGE_COMPONENT_WRITABLE_DIRECTORIES": "/opt/ml/output/data:/opt/ml/model:/opt/ml/input/data/ground_truth:/opt/ml/checkpoints:/tmp",
-                    "GRAND_CHALLENGE_COMPONENT_POST_CLEAN_DIRECTORIES": "/opt/ml/output/data:/opt/ml/model:/opt/ml/input/data/ground_truth",
                     "GRAND_CHALLENGE_COMPONENT_MAX_MEMORY_MB": "7168",
                     "GRAND_CHALLENGE_COMPONENT_SIGNING_KEY_HEX": "746f74616c6c79736563726574",
                 },
@@ -556,23 +556,23 @@ def test_handle_completed_job():
         time_limit=60,
         requires_gpu_type=GPUTypeChoices.NO_GPU,
         use_warm_pool=False,
-        signing_key=b"",
+        signing_key=b"itsasecret",
     )
 
-    return_code = 0
-
-    with io.BytesIO() as f:
-        f.write(
-            json.dumps(
-                {"return_code": return_code, "pk": f"algorithms-job-{pk}"}
-            ).encode("utf-8")
-        )
-        f.seek(0)
-        executor._s3_client.upload_fileobj(
-            Fileobj=f,
-            Bucket=settings.COMPONENTS_OUTPUT_BUCKET_NAME,
-            Key=executor._result_key,
-        )
+    content = json.dumps(
+        {"return_code": 0, "pk": f"algorithms-job-{pk}"}
+    ).encode("utf-8")
+    signature = hmac.new(
+        key=b"itsasecret", msg=content, digestmod=hashlib.sha256
+    ).hexdigest()
+    executor._s3_client.upload_fileobj(
+        Fileobj=io.BytesIO(content),
+        Bucket=settings.COMPONENTS_OUTPUT_BUCKET_NAME,
+        Key=executor._result_key,
+        ExtraArgs={
+            "Metadata": {"signature_hmac_sha256": signature},
+        },
+    )
 
     assert executor._handle_completed_job() is None