install

#!/usr/bin/env python2

import argparse
import yaml
import os
import subprocess
import shutil
import errno
import distutils.dir_util

def install_from_spack(spec, env):
    # Delete ~/.intel
    intel_home = os.path.join(os.environ['HOME'], '.intel')
    if os.path.isdir(intel_home):
        shutil.rmtree(intel_home)
    # Delete ~/intel
    intel_home = os.path.join(os.environ['HOME'], 'intel')
    if os.path.isdir(intel_home):
        shutil.rmtree(intel_home)
    # Delete cached intel license
    intel_lic = os.path.join(os.path.expanduser(env['spack_bin_root']), 'etc', 'spack', 'licenses', 'intel', 'license.lic')
    if os.path.isfile(intel_lic):
        os.remove(intel_lic)
    # Install the spec package if absent
    try:
        out = subprocess.check_output("spack find {}".format(spec), shell=True).decode("utf-8")
        if "No package" in out:
            print("{} not found. Build it.".format(spec))
            os.system("spack install --verbose {}".format(spec))
    except subprocess.CalledProcessError as e:
        pass

# Parse commandline options
parser = argparse.ArgumentParser(description='Deploy Spack software environment')

parser.add_argument('--env',
	type=argparse.FileType('r'),
	help='Environment file')
parser.add_argument('--install',
	action='store_true',
	help='Install packages')
parser.add_argument('--sync',
	action='store_true',
	help='Sync cached source files to the mirror site')
parser.add_argument('--refresh',
	action='store_true',
	help='Refresh Spack environment modules')

args = parser.parse_args()

# Read variables from configuration file
env = yaml.load(args.env)

# Raise exception if $spack_bin_root/bin/spack exists
if not os.path.isfile(os.path.join(os.path.expanduser(env['spack_bin_root']), 'bin', 'spack')):
	raise IOError(errno.ENOENT,
		os.strerror(errno.ENOENT),
		os.path.join(os.path.expanduser(env['spack_bin_root']), 'bin', 'spack'))

# Create ~/.spack directory if absent.
distutils.dir_util.mkpath(os.path.join(os.environ['HOME'], '.spack'))
# Create $spack_install_root if absent
distutils.dir_util.mkpath(env['spack_install_root'])
# Crate $spack_modules_root if absent
distutils.dir_util.mkpath(env['spack_modules_root'])

# Set and write to spack config file ~/.spack/config.yaml
if 'config' not in env.keys():
    env['config'] = {}
    env['config']['checksum'] = True
    env['config']['dirty'] = False
    env['config']['install_tree'] = {}
    env['config']['install_tree']['root'] = env['spack_install_root']
    env['config']['module_roots'] = {}
    env['config']['module_roots']['tcl'] = env['spack_modules_root']
    env['config']['source_cache'] = os.path.join(os.environ['HOME'], '.spack', 'sourcecache')
    env['config']['misc_cache'] = os.path.join(os.environ['HOME'], '.spack', 'misccache')
    env['config']['build_stage'] = ['$tempdir/$user']
    env['config']['db_lock_timeout'] = 180
out = {}
out['config'] = env['config']
with open(os.path.join(os.environ['HOME'], '.spack', 'config.yaml'), 'w') as outfile:
	yaml.dump(out, outfile, default_flow_style=False)

# Write ~/.spack/modules.yaml
if 'modules' not in env.keys():
    env['modules'] = {}
    env['modules']['tcl'] = {}
    env['modules']['tcl']['hash_length'] = '0'
    env['modules']['tcl']['naming_scheme'] = "'{name}/{version}-{compiler.name}-{compiler.version}'"
    env['modules']['tcl']['all'] = {}
    env['modules']['tcl']['all']['conflict'] = "'{name}'"
out = {}
out['modules'] = env['modules']
with open(os.path.join(os.environ['HOME'], '.spack', 'modules.yaml'), 'w') as outfile:
	yaml.dump(out, outfile, default_flow_style=False)

# Write ~/.spack/packages.yaml
out = {}
out['packages'] = env['packages']
if out['packages'] is not None:
    with open(os.path.join(os.environ['HOME'], '.spack', 'packages.yaml'), 'w') as outfile:
    	yaml.dump(out, outfile, default_flow_style=False)

# Write ~/.spack/compilers.yaml
out = {}
out['compilers'] = env['compilers']
with open(os.path.join(os.environ['HOME'], '.spack', 'compilers.yaml'), 'w') as outfile:
	yaml.dump(out, outfile, default_flow_style=False)

# Write ~/.spack/mirrors.yaml
out = {}
out['mirrors'] = env['mirrors']
with open(os.path.join(os.environ['HOME'], '.spack', 'mirrors.yaml'), 'w') as outfile:
	yaml.dump(out, outfile, default_flow_style=False)

# Install packages
if args.install:
 	for p in env['install']:
 		install_from_spack(p, env)

# Refresh TCL modulefiles
if args.refresh:
    os.system("spack module tcl refresh --delete-tree -y")

# Run Spack if --sync is NOT set
if args.sync:
    print("Sync cached source files to the mirror site")
    os.system("rsync -avr --progress %s/ root@proxy://var/lib/www/spack.pi.sjtu.edu.cn/mirror/" % os.path.join(os.environ['HOME'], '.spack', 'sourcecache'))