osg_annex_generate

#!/usr/bin/python3

import os
import sys
import getpass
import argparse
import re

queueName = "normal"
timeString = "01:00:00"
userName = getpass.getuser()
targetName = "stampede2-{0}.slurm".format(queueName)
annexName = "{0}@Stampede2-{1}".format(userName, queueName)
startExtra = 'Owner == \\"{0}\\"'.format(userName)

def verifyDuration(string):
	hhmmss = re.compile('^\d\d\:\d\d\:\d\d$')
	return hhmmss.match(string) is not None

parser = argparse.ArgumentParser(description="Generate a SLURM submit script which runs a pilot on Stampedede2.")
parser.add_argument("--queue",
	default=queueName,
	help="Specify the SLURM queue (partition); default {0}.".format(queueName))
parser.add_argument("--duration",
	default=timeString,
	help="Use hh:mmm:sss to specify the duration; default {0}.".format(timeString))
parser.add_argument("--target",
	default=targetName,
	help="Generate the SLURM script in this file; default {0}.".format(targetName))
parser.add_argument("--name",
	default=annexName,
	help="The annex name; default {0}.".format(annexName))

parser.add_argument("-N", "--nodes", default=1,
	help="The number of nodes to request, default 1.")
parser.add_argument("-n", "--ntasks", default=1,
	help="The number of tasks, default 1.")
parser.add_argument("-c", "--cpus-per-task", default=1,
	help="The number of CPUs per task, default 1.")

args = parser.parse_args()

queueName = args.queue
if verifyDuration(args.duration):
	timeString = args.duration
else:
	sys.exit("Duration '{0}' not in hh:mm:ss format, aborting.".format(args.duration))
targetName = args.target
annexName = args.name

header = '''
#SBATCH -J osgvo-pilot
#SBATCH -o osgvo-pilot/%j.out
#SBATCH -e osgvo-pilot/%j.err
'''

footer = '''
export GLIDEIN_Site="TACC"
export GLIDEIN_ResourceName="Stampede2"
export OSG_SQUID_LOCATION=""

module load tacc-singularity

singularity run --bind /cvmfs:/cvmfs docker://opensciencegrid/osgvo-docker-pilot:latest
'''

with open(targetName, 'w') as targetFile:
	print('#!/bin/bash', file=targetFile)
	targetFile.write(header)
	print('#SBATCH -p {0}'.format(queueName), file=targetFile)
	print('#SBATCH --nodes {0}'.format(args.nodes), file=targetFile)
	print('#SBATCH --ntasks {0}'.format(args.ntasks), file=targetFile)
	print('#SBATCH --cpus-per-task {0}'.format(args.cpus_per_task), file=targetFile)
	print('#SBATCH -t {0}'.format(timeString), file=targetFile)
	print('', file=targetFile)
	print('export ANNEX_NAME="{0}"'.format(annexName), file=targetFile)
	print('export TOKEN=`cat ~/token.txt`', file=targetFile)
	print('export GLIDEIN_Start_Extra="{0}"'.format(startExtra), file=targetFile)
	targetFile.write(footer)

os.chmod(targetName, 0o600)

print("SLURM script generated.  Perform the following steps to use it.")
print("    (1) Copy '{0}' to a Stampede2 login node.".format(targetName))
print("    (2) Copy your OSG Connect token (usually '~/token.txt') to '~/token.txt' on a Stampede2 login node.");
print("    (3) Create a directory named 'osgvo-pilot' in the directory to which you copied '{0}'.".format(targetName))
print("    (4) Run 'sbatch {0}'.".format(targetName))
sys.exit(0)