From fde10ce1ca8913c81d92130c7ee535f0053f6424 Mon Sep 17 00:00:00 2001 From: Sam Gillingham Date: Thu, 27 Jun 2024 08:34:51 +0000 Subject: [PATCH] AWS resets AWS_BATCH_JOB_ARRAY_INDEX so use a command line arg instead --- parallel_examples/awsbatch/do_prepare.py | 4 ++-- parallel_examples/awsbatch/do_tile.py | 17 +++++++++-------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/parallel_examples/awsbatch/do_prepare.py b/parallel_examples/awsbatch/do_prepare.py index 4a7bbab..df35f33 100755 --- a/parallel_examples/awsbatch/do_prepare.py +++ b/parallel_examples/awsbatch/do_prepare.py @@ -125,8 +125,8 @@ def main(): arrayProperties['size'] = len(colRowList) else: # must fake AWS_BATCH_JOB_ARRAY_INDEX - containerOverrides['environment'] = [{'name': 'AWS_BATCH_JOB_ARRAY_INDEX', - 'value': '0'}] + # can't set this as and env var as Batch overrides + containerOverrides['command'].extend(['--arrayindex', '0']) response = batch.submit_job(jobName="pyshepseg_tiles", jobQueue=cmdargs.jobqueue, diff --git a/parallel_examples/awsbatch/do_tile.py b/parallel_examples/awsbatch/do_tile.py index 9834623..affee72 100755 --- a/parallel_examples/awsbatch/do_tile.py +++ b/parallel_examples/awsbatch/do_tile.py @@ -21,13 +21,6 @@ gdal.UseExceptions() -# set by AWS Batch -ARRAY_INDEX = os.getenv('AWS_BATCH_JOB_ARRAY_INDEX') -if ARRAY_INDEX is None: - raise SystemExit('Must set AWS_BATCH_JOB_ARRAY_INDEX env var') - -ARRAY_INDEX = int(ARRAY_INDEX) - def getCmdargs(): """ @@ -48,9 +41,17 @@ def getCmdargs(): help="Maximum spectral difference for segmentation (default=%(default)s)") p.add_argument("--spectDistPcntile", type=int, default=50, required=False, help="Spectral Distance Percentile for segmentation (default=%(default)s)") + p.add_argument("--arrayindex", type=int, + help="Override AWS_BATCH_JOB_ARRAY_INDEX env var") cmdargs = p.parse_args() + if cmdargs.arrayindex is None: + cmdargs.arrayindex = os.getenv('AWS_BATCH_JOB_ARRAY_INDEX') + if cmdargs.arrayindex is None: + raise SystemExit('Must set AWS_BATCH_JOB_ARRAY_INDEX env var or ' + + 'specify --arrayindex') + return cmdargs @@ -75,7 +76,7 @@ def main(): tempDir = tempfile.mkdtemp() # work out which tile we are processing - col, row = dataFromPickle['colRowList'][ARRAY_INDEX] + col, row = dataFromPickle['colRowList'][cmdargs.arrayindex] # work out a filename to save with the output of this tile # Note: this filename format is repeated in do_stitch.py