Implement dynamic memory allocation for multi-frame AutoRIFT jobs

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [10.13.0]
 
 ### Added
+- Implemented dynamic memory allocation for multi-frame `AUTORIFT` jobs in `set_batch_overrides.py`. Sentinel-2 jobs with more than 4 frames now use 16 GB of memory instead of 8 GB. Single-frame Sentinel-2/Landsat and Sentinel-1 memory allcocation is unchanged.
 - Added optional `chip_size` and `search_range` parameters to the `ARIA_AUTORIFT.yml` job specification to enable user-defined `chip-size` and `search-range`.
 - Updated `AUTORIFT.yml`, `ARIA_AUTORIFT.yml`, and `ITS_LIVE_AUTORIFT.yml` validation schema to support processing of Sentinel-1D and Sentinel-2C/D granules.
 - Added a `model_context_length` parameter to the `OPERA_DIST_S1` job specification.

apps/set-batch-overrides/src/set_batch_overrides.py

@@ -1,9 +1,9 @@
 from math import ceil
 
 
-AUTORIFT_S2_MEMORY = '7875'
-AUTORIFT_LANDSAT_MEMORY = '15750'
-AUTORIFT_S1_MEMORY = '31500'
+AUTORIFT_MEMORY_8GB = '7875'
+AUTORIFT_MEMORY_16GB = '15750'
+AUTORIFT_MEMORY_32GB = '31500'
 
 RTC_GAMMA_10M_MEMORY = '63200'
 WATER_MAP_10M_MEMORY = '126000'
@@ -36,12 +36,18 @@ def get_granules(job_parameters: dict) -> list[str]:
 def get_autorift_memory(job_parameters: dict) -> str:
     granules = get_granules(job_parameters)
 
+    frame_counts = [len(job_parameters.get(k) or []) for k in ['granules', 'reference', 'secondary']]
+    num_frames = max(frame_counts)
+
     if granules[0].startswith('S2'):
-        return AUTORIFT_S2_MEMORY
+        if num_frames > 4:
+            return AUTORIFT_MEMORY_16GB
+        return AUTORIFT_MEMORY_8GB
+
     elif granules[0].startswith('L'):
-        return AUTORIFT_LANDSAT_MEMORY
+        return AUTORIFT_MEMORY_16GB
 
-    return AUTORIFT_S1_MEMORY
+    return AUTORIFT_MEMORY_32GB
 
 
 def get_insar_isce_burst_memory(job_parameters: dict) -> str:

tests/test_set_batch_overrides.py

@@ -1,9 +1,9 @@
 import pytest
 
 from set_batch_overrides import (
-    AUTORIFT_LANDSAT_MEMORY,
-    AUTORIFT_S1_MEMORY,
-    AUTORIFT_S2_MEMORY,
+    AUTORIFT_MEMORY_8GB,
+    AUTORIFT_MEMORY_16GB,
+    AUTORIFT_MEMORY_32GB,
     INSAR_ISCE_BURST_MEMORY_8G,
     INSAR_ISCE_BURST_MEMORY_16G,
     INSAR_ISCE_BURST_MEMORY_32G,
@@ -167,7 +167,7 @@ def test_set_batch_overrides_autorift_s1():
         'ResourceRequirements': [
             {
                 'Type': 'MEMORY',
-                'Value': AUTORIFT_S1_MEMORY,
+                'Value': AUTORIFT_MEMORY_32GB,
             }
         ],
         'Environment': [{'Name': 'OMP_NUM_THREADS', 'Value': '4'}],
@@ -185,12 +185,32 @@ def test_set_batch_overrides_autorift_s2():
         'ResourceRequirements': [
             {
                 'Type': 'MEMORY',
-                'Value': AUTORIFT_S2_MEMORY,
+                'Value': AUTORIFT_MEMORY_8GB,
             }
         ],
         'Environment': [{'Name': 'OMP_NUM_THREADS', 'Value': '1'}],
     }
 
+    assert lambda_handler(
+        {
+            'job_type': 'AUTORIFT',
+            'job_parameters': {
+                'granules': ['S2_stub'],  # Used for sensor detection
+                'reference': ['S2A_1', 'S2A_2', 'S2A_3', 'S2A_4', 'S2A_5'],
+                'secondary': ['S2B_1'],
+            },
+        },
+        None,
+    ) == {
+        'ResourceRequirements': [
+            {
+                'Type': 'MEMORY',
+                'Value': AUTORIFT_MEMORY_16GB,
+            }
+        ],
+        'Environment': [{'Name': 'OMP_NUM_THREADS', 'Value': '2'}],
+    }
+
 
 def test_set_batch_overrides_autorift_landsat():
     assert lambda_handler(
@@ -203,7 +223,7 @@ def test_set_batch_overrides_autorift_landsat():
         'ResourceRequirements': [
             {
                 'Type': 'MEMORY',
-                'Value': AUTORIFT_LANDSAT_MEMORY,
+                'Value': AUTORIFT_MEMORY_16GB,
             }
         ],
         'Environment': [{'Name': 'OMP_NUM_THREADS', 'Value': '2'}],