#!/usr/bin/env python3

# This is the launcher to be used by mpich
# to run jobs on remote notes.

from subprocess import *
import sys
import os
import re
image = os.environ["SINGULARITY_CONTAINER"]
host = sys.argv[1:2]
if host == ["-x"]:
    host = sys.argv[2:3]
    cmd = sys.argv[3:]
else:
    cmd = sys.argv[2:]
opts = re.split(r'\s+',os.environ["SING_OPTS"].strip())

# We need to make sure we export our environment over MPI
envargs = ["env", "-i"]
for key in os.environ:
    # Keep slurm variables out.
    if "SLURM" not in key and re.match(r'^\w+$', key):
        value = os.environ[key]
        if re.search(r'[\[\(\*\s]', value):
            if not re.search(r"'", value):
                envargs += ["%s='%s'" % (key, value)]
            elif not re.search(r'"', value):
                envargs += ['%s="%s"' % (key, value)]
            else:
                print("Warn: not including variable: '%s'" % key, file=sys.stderr)
        else:
            envargs += ["%s=%s" % (key, value)]

full_cmd = ["ssh"] + host + envargs + ["singularity","exec"] + opts +[image] + cmd

call(full_cmd) 
