array_worker.py

#!/bin/env python

from argparse import ArgumentParser
from glob import glob

from slurm_utils import Chunk, Context

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("data_dir", help="Directory containing data files to process")
    args = parser.parse_args()

    data_files = glob(f"{args.data_dir}/*")

    with Context( verbosity=2 ) as context:

        if( context.is_array() ):
            context.log.debug(f"Running in a SLURM array")
        else:
            pass
        
        with Chunk( data_files, context=context ) as chunk:
            for file in chunk:
                with context.mkTempFile() as f:
                    with open(file) as input:
                        for line in input:
                            f.write(line.upper())
                
                # Do stuff...
        
        context.collect()