large-mem.py

#!/usr/bin/env python3
import argparse
import gc
import numpy as np
import os
import pathlib
import psutil


kB = 1024
MB = 1024 * kB

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("file", type=pathlib.Path)
  parser.add_argument("-c","--chunk-size", type=int, default=10*MB)
  args = parser.parse_args()
  
  proc  = psutil.Process()
  data = np.memmap(args.file, dtype='uint8', mode='r')


  i = 0
  while True:
      j = i + args.chunk_size
      chunk = data[i: j]
      i = j
      
      # Do work
      chunk.sum()
      gc.collect()

      print(f"Memory: {proc.memory_full_info().rss} MB")