feat: delete s2 granules after success (#309)

Author: ceholden

lambda_functions/cleanup_sentinel2_granules.py

@@ -0,0 +1,75 @@
+"""
+HLS: Remove Sentinel-2 granules after successful processing
+
+This Lambda function is connected in our workflow after the check for
+failures in the "success" pathway. As such we do not need to ensure the
+granule has been successfully processed within this scope.
+
+The only complication involved is handling of "twin" granules because
+we need to ensure the first granule processing workflow does not delete
+the 2nd of the "twin" granules. A "twin" granule occurs when the same
+tile and date are acquired twice due to the satellite switching receiving
+stations during the downlink. Two granules exist for the same date and tile,
+and for complete coverage we must process and combine the two granules
+together.
+
+A "twin" granule is processed in two steps,
+
+1. The first of the "twin" is downloaded into the input bucket, triggering a
+   workflow to process this granule.
+2. This first workflow only finds one of the two granules that will eventually be
+   downloaded, and only processes the first.
+3. The second of the "twin" granules is downloaded into the input bucket, triggering
+   a workflow to process this granule.
+4. This second workflow finds two input granule IDs, and processes them both.
+
+In this scenario we only want to delete the inputs for the second workflow.
+"""
+
+import os
+
+import boto3
+
+
+s3 = boto3.client("s3")
+bucket = os.getenv("SENTINEL_INPUT_BUCKET", None)
+if bucket is None:
+    raise Exception("No Input Bucket set")
+
+
+def handler(event: dict, context: dict):
+    # We may receive 2 granules split by a comma if this is a twin granule workflow
+    granules = event["granule"].split(",")
+
+    prefixes = {granule[0:-6] for granule in granules}
+    if len(prefixes) != 1:
+        raise ValueError(f"Received {len(prefixes)} granule prefixes")
+    prefix = list(prefixes)[0]
+
+    response = s3.list_objects_v2(
+        Bucket=bucket,
+        Prefix=prefix,
+    )
+
+    granule_zips = [obj["Key"] for obj in response["Contents"]]
+
+    # We have three possible cases,
+    # 1. Non-twin granule (1 ID, 1 zip)
+    # 2. Twin granule input and twin granule job (2 IDs, 2 zips)
+    if len(granules) == len(granule_zips):
+        if len(granules) > 1:
+            print(f"Deleting inputs of twin granule case: {granule_zips}")
+        else:
+            print(f"Deleting input of single granule case: {granule_zips[0]}")
+
+        for granule_zip in granule_zips:
+            s3.delete_object(Bucket=bucket, Key=granule_zip)
+        return granule_zips
+
+    # 3. Twin granules downloaded but only one was processed in this workflow
+    else:
+        print(
+            "Twin granule case detected but this workflow did not process it. "
+            f"Skipping deletion (IDs={granules}, zips={granule_zips})"
+        )
+        return []

lambda_functions/tests/test_cleanup_sentinel2_granules.py

@@ -0,0 +1,79 @@
+from unittest.mock import patch
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def env_setup(monkeypatch):
+    monkeypatch.setenv("SENTINEL_INPUT_BUCKET", "sentinelinput")
+
+
+@pytest.fixture
+def mock_delete_object():
+    from lambda_functions.cleanup_sentinel2_granules import s3
+
+    with patch.object(s3, "delete_object") as mock_delete_object:
+        yield mock_delete_object
+
+
+def test_single_granule_case(mock_delete_object, capsys):
+    from lambda_functions.cleanup_sentinel2_granules import handler, s3
+
+    with patch.object(
+        s3,
+        "list_objects_v2",
+        return_value={"Contents": [{"Key": "one.zip"}]},
+    ):
+        handler({"granule": "one12345"}, {})
+        assert mock_delete_object.call_count == 1
+
+    captured = capsys.readouterr()
+    assert "Deleting input of single granule case" in captured.out
+
+
+def test_twin_granule_only_one_skip_delete(mock_delete_object, capsys):
+    from lambda_functions.cleanup_sentinel2_granules import handler, s3
+
+    with patch.object(
+        s3,
+        "list_objects_v2",
+        return_value={"Contents": [{"Key": "one.zip"}, {"Key": "two.zip"}]},
+    ):
+        handler({"granule": "one12345"}, {})
+        mock_delete_object.assert_not_called()
+
+    captured = capsys.readouterr()
+    assert (
+        "Twin granule case detected but this workflow did not process it. Skipping"
+        in captured.out
+    )
+
+
+def test_twin_granule_has_both_deletes_both(mock_delete_object, capsys):
+    from lambda_functions.cleanup_sentinel2_granules import handler, s3
+
+    with patch.object(
+        s3,
+        "list_objects_v2",
+        return_value={"Contents": [{"Key": "one.zip"}, {"Key": "two.zip"}]},
+    ):
+        handler({"granule": "one12345_111111,one12345_222222"}, {})
+        assert mock_delete_object.call_count == 2
+
+    captured = capsys.readouterr()
+    assert "Deleting inputs of twin granule case" in captured.out
+
+
+def test_bad_granule_id_inputs(mock_delete_object):
+    from lambda_functions.cleanup_sentinel2_granules import handler, s3
+
+    with (
+        patch.object(
+            s3,
+            "list_objects_v2",
+            return_value={"Contents": [{"Key": "one.zip"}, {"Key": "two.zip"}]},
+        ) as mock_list_objects,
+        pytest.raises(ValueError, match=r"Received 2 granule prefixes"),
+    ):
+        handler({"granule": "thisiswronggranuleid,obviouslynotagranuleid"}, {})
+        mock_list_objects.assert_not_called()
+        mock_delete_object.assert_not_called()

stack/hlsconstructs/lambdafunc.py

@@ -27,7 +27,7 @@ def __init__(
         code_str: str = None,
         package_code_dir: str = None,
         env: Dict = None,
-        runtime: aws_lambda.Runtime = aws_lambda.Runtime.PYTHON_3_8,
+        runtime: aws_lambda.Runtime = aws_lambda.Runtime.PYTHON_3_9,
         handler: str = "index.handler",
         layers: list = None,
         cron_str: str = None,

stack/hlsconstructs/sentinel_step_function.py

@@ -1,7 +1,7 @@
 import json
-from typing import Union
+from typing import Optional, Union
 
-from aws_cdk import aws_iam, aws_stepfunctions
+from aws_cdk import aws_stepfunctions
 from constructs import Construct
 from hlsconstructs.batch_step_function import BatchStepFunction
 from hlsconstructs.lambdafunc import Lambda
@@ -22,6 +22,7 @@ def __init__(
         sentinel_ac_logger: Lambda,
         sentinel_logger: Lambda,
         check_exit_code: Lambda,
+        cleanup_granule: Optional[Lambda],
         replace_existing: bool,
         gibs_outputbucket: str,
         debug_bucket: Union[bool, str] = False,
@@ -124,24 +125,27 @@ def __init__(
                     "Type": "Task",
                     "Resource": sentinel_ac_logger.function.function_arn,
                     "Next": "CheckSentinelExitCode",
+                    "ResultPath": "$.exitCode",
                     "Retry": [retry],
                 },
                 "CheckSentinelExitCode": {
                     "Type": "Task",
                     "Resource": check_exit_code.function.function_arn,
                     "Next": "HadSentinelFailure",
+                    "InputPath": "$.exitCode",
+                    "ResultPath": "$.success",
                     "Retry": [retry],
                 },
                 "HadSentinelFailure": {
                     "Type": "Choice",
                     "Choices": [
                         {
-                            "Variable": "$",
+                            "Variable": "$.success",
                             "BooleanEquals": True,
                             "Next": "Done",
                         },
                         {
-                            "Variable": "$",
+                            "Variable": "$.success",
                             "BooleanEquals": False,
                             "Next": "Error",
                         },
@@ -153,6 +157,20 @@ def __init__(
             },
         }
 
+        # Add "cleanup" step to delete successfully processed granules for
+        # forward processing, but not historic
+        if cleanup_granule is not None:
+            sentinel_state_definition["States"]["CleanupGranule"] = {
+                "Type": "Task",
+                "Resource": cleanup_granule.function.function_arn,
+                "Next": "Done",
+                "InputPath": "$",
+                "Retry": [retry],
+            }
+            sentinel_state_definition["States"]["HadSentinelFailure"]["Choices"][0][
+                "Next"
+            ] = "CleanupGranule"
+
         if debug_bucket:
             sentinel_state_definition["States"]["ProcessSentinel"]["Parameters"][
                 "ContainerOverrides"

stack/stack.py

@@ -304,7 +304,7 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
                     os.path.dirname(__file__), "..", "layers", "hls_lambda_layer"
                 )
             ),
-            compatible_runtimes=[aws_lambda.Runtime.PYTHON_3_8],
+            compatible_runtimes=[aws_lambda.Runtime.PYTHON_3_9],
         )
 
         self.pr2mgrs_lambda = Lambda(
@@ -492,6 +492,14 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             layers=[self.hls_lambda_layer],
         )
 
+        self.cleanup_sentinel2_granule = Lambda(
+            self,
+            "CleanupSentinelSuccesses",
+            code_file="cleanup_sentinel2_granules.py",
+            env={"SENTINEL_INPUT_BUCKET": SENTINEL_INPUT_BUCKET},
+            timeout=120,
+        )
+
         self.get_random_wait = Lambda(
             self,
             "GetRandomWait",
@@ -625,6 +633,7 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             sentinel_ac_logger=self.sentinel_ac_logger,
             sentinel_logger=self.sentinel_logger,
             check_exit_code=self.check_exit_code,
+            cleanup_granule=self.cleanup_sentinel2_granule,
             outputbucket_role_arn=OUTPUT_BUCKET_ROLE_ARN,
             replace_existing=REPLACE_EXISTING,
             gibs_outputbucket=GIBS_OUTPUT_BUCKET,
@@ -643,6 +652,9 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             sentinel_ac_logger=self.sentinel_ac_logger,
             sentinel_logger=self.sentinel_logger_historic,
             check_exit_code=self.check_exit_code,
+            # Do not cleanup granules for historic workflow to avoid
+            # twin granule race condition
+            cleanup_granule=None,
             outputbucket_role_arn=OUTPUT_BUCKET_ROLE_ARN,
             replace_existing=REPLACE_EXISTING,
             gibs_outputbucket=GIBS_OUTPUT_BUCKET_HISTORIC,
@@ -1042,6 +1054,20 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             self.sentinel_input_bucket_historic_policy
         )
 
+        self.cleanup_sentinel_input_bucket_policy = aws_iam.PolicyStatement(
+            resources=[
+                self.sentinel_input_bucket.bucket_arn,
+                f"{self.sentinel_input_bucket.bucket_arn}/*",
+            ],
+            actions=[
+                "s3:List*",
+                "s3:DeleteObject",
+            ],
+        )
+        self.cleanup_sentinel2_granule.function.add_to_role_policy(
+            self.cleanup_sentinel_input_bucket_policy
+        )
+
         self.laads_task.role.add_to_policy(
             aws_iam.PolicyStatement(
                 resources=[