<commit_before>#!/usr/bin/env python
# Metarunlog, experiment management tool.
# Author: <NAME>
# Date: 2015-01-23

import metarunlog.cfg as cfg # NOTE cfg is modified by MetaRunLog._loadBasedirConfig() with custom configuration.
import os
import sys
from os import listdir
from os.path import isdir, isfile, join, relpath, expanduser
import argparse
import subprocess
import json
import pdb
import datetime
from collections import OrderedDict
from shutil import copy as shcopy
from jinja2 import Template, Environment, meta
import itertools
import getpass

DEBUG = True

class InvalidOutDirException(Exception):
    pass
class NoCleanStateException(Exception):
    pass
class NoBasedirConfig(Exception):
    def __init__(self, basedir, e):
        self.basedir = basedir
        self.orig_e = e
class InvalidExpIdException(Exception):
    pass
class BatchException(Exception):
    pass
class HpcException(Exception):
    pass
class NoSuchCustomCommandException(Exception):
    pass

def get_commit():
    cline = subprocess.check_output("git log -n1 --oneline", shell=True)
    #print "cline: ", cline
    cline = cline.split()
    return (cline[0], " ".join(cline[1:]))

def initBasedir(basedir, args):
    """ init basedir"""
    # check if it already is a valid basedir, then raise an exception
    try:
        mrlState = MetaRunLog(basedir)
    except NoBasedirConfig:
        pass # should be raised
    else:
        return "{} already has valid .mrl.cfg file, remove it first to re-init.".format(basedir)
    # initialize .mrl.cfg file as a json-dict copy of cfg.py template
    with open(join(basedir, '.mrl.cfg'),"w") as fh:
        # copy the cfg attributes into ordered dict
        bconf = {k:getattr(cfg,k) for k in dir(cfg) if '__' not in k}
        bconf['outdir'] = args.outdir
        json.dump(bconf, fh, indent=2)
        fh.write("\n")
    try:
        mrlState = MetaRunLog(basedir)
    except Exception as e:
        print("initBasedir() - Failed to init basedirconfig")
        raise
    else:
        return join(basedir, ".mrl.cfg")

class BatchParser:
    """
    BatchParser class is initialized with batch-template (list of lines)
    including the last lines determining the values,
    and generates self.output as a list of strings (that can be written to file).
    Raises xxx.
    """
    def __init__(self, template):
        grid = OrderedDict()
        params = [{}]
        # iterate over last lines and parse them to collect grid parameters.
        # correct formatting of mrl instructions:
        # [something] MRL:grid['key'] = [vallist]
        # where [something] is comment symbol. The space after it is essential.
        for line in template[::-1]:
            if not line.strip(): continue
            line = line.split(None, 1)[-1].split(":",1) #discard comment symbol
            if line[0] != 'MRL': break
            exec(line[1])
        # construct output params list
        keys = grid.keys()
        vals = list(itertools.product(*[grid[k] for k in keys]))
        self.params = [dict(zip(keys, v)) for v in vals]
        self.params = [dict(d1.items() + d2.items()) for d1,d2 in itertools.product(self.params, params)]
        # render the templates
        self.output = []
        jtmpl = Template("".join(template))
        for i, param in enumerate(self.params):
            self.output.append((i+1, param, jtmpl.render(**param)))

class MetaRunLog:
    """
    Metarunlog state from which all actions are coordinated.
    Needs to be run from a valid basedir (with .mrl.cfg file).
    Initialization overwrites cfg module from .mrl.cfg: metarunlog.cfg.py should be
    seen as a template for .mrl.cfg.
    """
    def __init__(self, basedir):
        self.basedir = basedir
        self._loadBasedirConfig()
        self.outdir  = join(self.basedir, cfg.outdir)
        if not isdir(self.outdir): raise InvalidOutDirException("Need output directory " + self.outdir + "  Fix your .mrl.cfg file")
        self.expDirList = sorted([x for x in listdir(self.outdir) if self._checkValidExp(x)])
        self.expList = [int(x) for x in self.expDirList]
        self.lastExpId = None if not self.expList else self.expList[-1]
        self.hpcPass = None
        # parse the custom commands into templates and extract the optional arguments for argparse
        self.jEnv = Environment()
        self.customTemplates = {customName:[]  for customName in cfg.custom.keys()}
        self.customOptVars = {customName:set() for customName in cfg.custom.keys()}
        for customName, cList in cfg.custom.iteritems():
            for cmdEnv in cList:
                self.customTemplates[customName].append({envKey: self.jEnv.from_string(cmd) for envKey, cmd in cmdEnv.iteritems()})
                for cmd in cmdEnv.values():
                    self.customOptVars[customName].update(meta.find_undeclared_variables(self.jEnv.parse(cmd)))

    def _loadBasedirConfig(self):
        try:
            with open(join(self.basedir, '.mrl.cfg')) as fh:
                bconf = json.load(fh)
                for k,v in bconf.iteritems():
                    setattr(cfg,k,v)
            return True
        except Exception as e:
            raise NoBasedirConfig(self.basedir, str(e))

    def new(self,args):
        expId = 1 if not self.lastExpId else self.lastExpId+1
        expConfig = OrderedDict()
        expConfig ['expId'] = expId
        expConfig['basedir'] = self.basedir
        gitclean = not bool(args.notclean)
        expConfig['gitFailUntracked']= args.gitFailUntracked
        untracked = 'no' if args.gitFailUntracked == 'no' else 'normal'
        uncommited = subprocess.check_output("git status --porcelain --untracked=" + untracked, shell=True)
        #print "uncommited: ", uncommited
        if gitclean and uncommited:
            raise NoCleanStateException("new: uncommited files -- please commit changes first\n" + uncommited)
        expConfig['gitHash'], expConfig['gitDescription'] = get_commit()
        if not gitclean and uncommited: expConfig['gitHash'] += '-sloppy'
        expConfig['timestamp'] = datetime.datetime.now().isoformat().split('.')[0]
        expConfig['user'] = getpass.getuser()
        expConfig['shortDescription'] = args.description if args.description else ""
        expConfig['longDescription'] = "" # edit into .mrl file directly
        # make dir and symlink  and
        expDir = self._getExpDir(expId, True)
        if self.lastExpId: os.remove(join(self.outdir,'last'))
        os.symlink(relpath(expDir,self.outdir), join(self.outdir, 'last'))
        os.mkdir(expDir)
        # After this point: expDir is made, no more exception throwing! copy config files
        try:
            if args.copyConfigFrom == 'no':
                pass
            else:
                srcDir = self._getExpDir(self._resolveExpId(args.copyConfigFrom))
                self._copyConfigFrom(srcDir, expDir)
        except (InvalidExpIdException, IOError) as e:
            print("Can't copy config files. ")
            print(e)
            print("Still succesfully created new experiment directory.")
        self._saveExpDotmrl(expDir, expConfig)
        self._renderLaunchScript(expDir, expConfig, self._fmtLaunchScriptFile(expId))
        self.lastExpId = expId
        return expDir

    def info(self, args):
        """ load info from experiment id and print it """
        expId = self._resolveExpId(args.expId)
        expConfig = self._getExpConf(expId)
        return "\n".join("{} : {}".format(k,v) for k,v in expConfig.iteritems())

    def ls(self, args):
        for expId in self.expList[::-1]:
            expConfig = self._getExpConf(expId)
            row = self._fmtSingleExp(expId)
            if args.tm:
                row += "\t" + expConfig['timestamp']
            if args.ghash:
                row += "\t" + expConfig['gitHash']
            if args.gdesc:
                row += "\t" + expConfig['gitDescription']
            if args.desc:
                row += "\t" + expConfig['shortDescription']
            print row
        return ""

    def batch(self, args):
        # resolve id
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        expConfig = self._getExpConf(expId)
        # check if already expanded in batch and cancel
        oldBatchList = []
        if 'batchlist' in expConfig:
            if args.replace:
                oldBatchList = expConfig['batchlist']
            else:
                raise BatchException("Experiment {} is already expanded into a batch of length {}".\
                        format(expId, len(expConfig['batchlist'])))
        # make BatchParser object
        with open(join(expDir, cfg.batchTemplFile)) as fh:
            bp = BatchParser(fh.readlines())
        # check if BatchParser output is non-empty
        if not bp.output:
            err = "BatchParser output is empty, are you sure {} is a batch template?"
            err.format(join(expDir, cfg.batchTemplFile))
            raise BatchException(err)
        # generate output directories, write the output config files
        expConfig['batchlist'] = []
        for i, params, fileContent in bp.output:
            expConfig['batchlist'].append(params)
            subExpDir = join(expDir, self._fmtBatchExp(expId=expId, subExpId=i))
            if i > len(oldBatchList):
                os.mkdir(subExpDir)
            with open(join(subExpDir, cfg.batchTemplFile), "w") as fh:
                fh.write(fileContent)
                fh.write("\n")
            self._renderLaunchScript(subExpDir, expConfig, self._fmtLaunchScriptFile(expId, i))
        # update the current .mrl file
        self._saveExpDotmrl(expDir, expConfig)
        self._writeqsubFile(expId, len(bp.output))
        os.remove(join(expDir, self._fmtLaunchScriptFile(expId))) # this conf is a template so doesnt make sense to run it
        return "Succesfully generated config files for expId {}.\n{}".format(expId,expConfig['batchlist'])

    def hpcSubmit(self, args):
        # TODO skip the qsub file and qsub all from mrl. Use _getRunLocations function
        # TODO Use custom() for this? Add pwd function
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        remoteExpDir = join(cfg.hpcBasedir, cfg.outdir, self._fmtSingleExp(expId))
        expConfig = self._getExpConf(expId)
        if 'hpcSubmit' in expConfig and not args.replace:
            raise HpcException("hpcSubmit key already in {} .mrl file. Aborting to avoid data loss.".format(expId))
        # scp to copy to remote
        try:
            self._hpc("mkdir {}".format(remoteExpDir))
        except Exception as e:
            if not args.replace:
                raise HpcException("hpc submission - error in mkdir returned nonzero exit. Dir already exists?")
        self._hpc("scp -r {} {}:{}".format(join(expDir,'*'), cfg.hpcServer, remoteExpDir), ssh=False)
        # subprocess ssh then qsub then print qstat
        if 'batchlist' in expConfig:
            hpcRet = self._hpc('cd {}; ./{}'.format(remoteExpDir, cfg.qsubFile))
            print hpcRet.strip().split("\n")
            # every even-indexed line is an id
            try:
                qsubids = [int(x) for i,x in enumerate(hpcRet.strip().split("\n")) if i%2==0]
            except Exception as e: #failure on qsub? save output
                print "Couldn't save qsubids correctly: ", e
                qsubids = hpcRet
        else:
            hpcRet = self._hpc('cd {}; qsub {}'.format(remoteExpDir, self._fmtLaunchScriptFile(expId)))
            print hpcRet
            qsubids = int(hpcRet)
        expConfig['hpcSubmit'] = qsubids
        self._saveExpDotmrl(expDir, expConfig)

    def hpcFetch(self, args):
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        remoteExpDir = join(cfg.hpcBasedir, cfg.outdir, self._fmtSingleExp(expId))
        expConfig = self._getExpConf(expId)
        if 'hpcSubmit' not in expConfig:
            raise HpcException("hpcSubmit key not in in {} .mrl file. Was this launched as hpc job?".format(expId))
        # Locations to fetch from
        runLocs = self._getRunLocations(expId, args.subExpId, expConfig, relativeTo=expDir)
        # Exclude datafiles ifneeded
        if not args.data:
            excludes = " ".join(['--exclude={}'.format(datafile) for datafile in cfg.hpcOutputData])
        else:
            excludes = ""
        for loc in runLocs:
            cmd = "rsync -aP {} {}:{}/ {}/".format(excludes, cfg.hpcServer, join(remoteExpDir,loc), join(expDir,loc))
            self._hpc(cmd, ssh=False)

    def custom(self, args):
        customName = args.customName
        cmdTemplates = self.customTemplates[customName]
        expId = self._resolveExpId(args.expId)
        expConfig = self._getExpConf(expId)
        runLocs = self._getRunLocations(expId, args.subExpId, expConfig, relativeTo=self.outdir)
        originalCwd = os.getcwd() # self.basedir if ran through ./mrl
        cmdParams = {'mrlOutdir': self.outdir, 'mrlBasedir': self.basedir}
        cmdParams.update(expConfig)
        cmdParams.update({k:v for k,v in vars(args).items() if v}) # the optional params if supplied
        # Render & run the commands for all experiments
        for relloc in runLocs:
            absloc = join(self.outdir, relloc)
            print "{} for location {}".format(customName, relloc)
            # Make location cmdParams
            cmdParams.update({'relloc': relloc, 'absloc': absloc})
            # Execute commands in their environment
            for cmdEnvTemplate in cmdTemplates:
                cmdEnv = {k:tmpl.render(cmdParams) for k,tmpl in cmdEnvTemplate.items()}
                if 'cwd' in cmdEnv:
                    os.chdir(os.path.expanduser(cmdEnv['cwd']))
                print cmdEnv['command']
                stdout = open(join(absloc, cmdEnv['output']),"w") if 'output' in cmdEnv else None
                subprocess.call(cmdEnv['command'], stdout=stdout, shell=True)
                if 'cwd' in cmdEnv:
                    os.chdir(originalCwd)

    def _hpc(self, cmd, ssh=True):
        if not self.hpcPass:
            self.hpcPass= getpass.getpass("Password for hpc: ")
        if ssh:
            sshcmd = "ssh {}".format(cfg.hpcServer)
            cmd = '"{}"'.format(cmd.replace('"', '\\"'))
        else:
            sshcmd = ""
        mcmd = "sshpass -p '{}' {} {}".format(self.hpcPass, sshcmd, cmd)
        cleancmd = "sshpass -p '{}' {} {}".format('***', sshcmd, cmd)
        if DEBUG: print cleancmd
        try:
            out = subprocess.check_output(mcmd, shell=True)
        except subprocess.CalledProcessError as e:
            # remove pwd from cmd
            err = str(e).replace(mcmd, cleancmd)
            raise HpcException(err)
        return out

    def _getRunLocations(self, expId, subExpId, expConfig, relativeTo=''):
        if 'batchlist' in expConfig:
            if subExpId == 'all':
                locs = [self._fmtBatchExp(expId, i) for i in range(1,len(expConfig['batchlist'])+1)]
            elif subExpId.isdigit():
                if int(subExpId) > len(expConfig['batchlist']):
                    raise SubExpIdException("subExpId {} out of range (batch size {} in expConfig)".format(subExpId,len(expConfig['batchlist'])))
                locs = [self._fmtBatchExp(expId, int(subExpId)), ]
            else:
                # TODO list of subexpids can be handy
                raise SubExpIdException("Don't understand subExpId {}.".format(subExpId))
        else:
            locs = ['']
        expDir = self._getExpDir(expId)
        if relativeTo != expDir:
            locs = [join(expDir, loc) for loc in locs]
            if relativeTo:
                locs = [relpath(loc, relativeTo) for loc in locs]
        return locs

    def _writeqsubFile(self, expId, N):
        # TODO make git clone  and checkout part of this to avoid 100 clones to the subdirs.
        expDir = self._getExpDir(expId)
        command = 'cd {} && qsub {} && cd .. && sleep 1 && echo "qsubbed {}"'
        subdirs = [self._fmtBatchExp(expId=expId, subExpId=i) for i in range(1,N+1)]
        lsfiles = [self._fmtLaunchScriptFile(expId, i) for i in range(1,N+1)]
        cmds = [command.format(subdir, lsfile, subdir) for lsfile,subdir in zip(lsfiles,subdirs)]
        cmds.append("")
        open(join(expDir, cfg.qsubFile),"w").write("\n".join(cmds))
        os.chmod(join(expDir, cfg.qsubFile), 0770)

    def _renderLaunchScript(self, expDir, expConfig, fn):
        """ Launch script is jinja template in basedir config.
        It has access to expDir, the global configuration,
        and all experiment specific variables
        in expConfig."""
        # TODO make git clone and checkout part optional here
        jtemp = Template(cfg.launchScriptTempl)
        tparams = {k:getattr(cfg, k) for k in dir(cfg) if '_' not in k}
        tparams.update(expConfig)<commit_msg>copy all configfiles to batch subdir<commit_after>#!/usr/bin/env python
# Metarunlog, experiment management tool.
# Author: <NAME>
# Date: 2015-01-23

import metarunlog.cfg as cfg # NOTE cfg is modified by MetaRunLog._loadBasedirConfig() with custom configuration.
import os
import sys
from os import listdir
from os.path import isdir, isfile, join, relpath, expanduser
import argparse
import subprocess
import json
import pdb
import datetime
from collections import OrderedDict
from shutil import copy as shcopy
from jinja2 import Template, Environment, meta
import itertools
import getpass

DEBUG = True

class InvalidOutDirException(Exception):
    pass
class NoCleanStateException(Exception):
    pass
class NoBasedirConfig(Exception):
    def __init__(self, basedir, e):
        self.basedir = basedir
        self.orig_e = e
class InvalidExpIdException(Exception):
    pass
class BatchException(Exception):
    pass
class HpcException(Exception):
    pass
class NoSuchCustomCommandException(Exception):
    pass

def get_commit():
    cline = subprocess.check_output("git log -n1 --oneline", shell=True)
    #print "cline: ", cline
    cline = cline.split()
    return (cline[0], " ".join(cline[1:]))

def initBasedir(basedir, args):
    """ init basedir"""
    # check if it already is a valid basedir, then raise an exception
    try:
        mrlState = MetaRunLog(basedir)
    except NoBasedirConfig:
        pass # should be raised
    else:
        return "{} already has valid .mrl.cfg file, remove it first to re-init.".format(basedir)
    # initialize .mrl.cfg file as a json-dict copy of cfg.py template
    with open(join(basedir, '.mrl.cfg'),"w") as fh:
        # copy the cfg attributes into ordered dict
        bconf = {k:getattr(cfg,k) for k in dir(cfg) if '__' not in k}
        bconf['outdir'] = args.outdir
        json.dump(bconf, fh, indent=2)
        fh.write("\n")
    try:
        mrlState = MetaRunLog(basedir)
    except Exception as e:
        print("initBasedir() - Failed to init basedirconfig")
        raise
    else:
        return join(basedir, ".mrl.cfg")

class BatchParser:
    """
    BatchParser class is initialized with batch-template (list of lines)
    including the last lines determining the values,
    and generates self.output as a list of strings (that can be written to file).
    Raises xxx.
    """
    def __init__(self, template):
        grid = OrderedDict()
        params = [{}]
        # iterate over last lines and parse them to collect grid parameters.
        # correct formatting of mrl instructions:
        # [something] MRL:grid['key'] = [vallist]
        # where [something] is comment symbol. The space after it is essential.
        for line in template[::-1]:
            if not line.strip(): continue
            line = line.split(None, 1)[-1].split(":",1) #discard comment symbol
            if line[0] != 'MRL': break
            exec(line[1])
        # construct output params list
        keys = grid.keys()
        vals = list(itertools.product(*[grid[k] for k in keys]))
        self.params = [dict(zip(keys, v)) for v in vals]
        self.params = [dict(d1.items() + d2.items()) for d1,d2 in itertools.product(self.params, params)]
        # render the templates
        self.output = []
        jtmpl = Template("".join(template))
        for i, param in enumerate(self.params):
            self.output.append((i+1, param, jtmpl.render(**param)))

class MetaRunLog:
    """
    Metarunlog state from which all actions are coordinated.
    Needs to be run from a valid basedir (with .mrl.cfg file).
    Initialization overwrites cfg module from .mrl.cfg: metarunlog.cfg.py should be
    seen as a template for .mrl.cfg.
    """
    def __init__(self, basedir):
        self.basedir = basedir
        self._loadBasedirConfig()
        self.outdir  = join(self.basedir, cfg.outdir)
        if not isdir(self.outdir): raise InvalidOutDirException("Need output directory " + self.outdir + "  Fix your .mrl.cfg file")
        self.expDirList = sorted([x for x in listdir(self.outdir) if self._checkValidExp(x)])
        self.expList = [int(x) for x in self.expDirList]
        self.lastExpId = None if not self.expList else self.expList[-1]
        self.hpcPass = None
        # parse the custom commands into templates and extract the optional arguments for argparse
        self.jEnv = Environment()
        self.customTemplates = {customName:[]  for customName in cfg.custom.keys()}
        self.customOptVars = {customName:set() for customName in cfg.custom.keys()}
        for customName, cList in cfg.custom.iteritems():
            for cmdEnv in cList:
                self.customTemplates[customName].append({envKey: self.jEnv.from_string(cmd) for envKey, cmd in cmdEnv.iteritems()})
                for cmd in cmdEnv.values():
                    self.customOptVars[customName].update(meta.find_undeclared_variables(self.jEnv.parse(cmd)))

    def _loadBasedirConfig(self):
        try:
            with open(join(self.basedir, '.mrl.cfg')) as fh:
                bconf = json.load(fh)
                for k,v in bconf.iteritems():
                    setattr(cfg,k,v)
            return True
        except Exception as e:
            raise NoBasedirConfig(self.basedir, str(e))

    def new(self,args):
        expId = 1 if not self.lastExpId else self.lastExpId+1
        expConfig = OrderedDict()
        expConfig ['expId'] = expId
        expConfig['basedir'] = self.basedir
        gitclean = not bool(args.notclean)
        expConfig['gitFailUntracked']= args.gitFailUntracked
        untracked = 'no' if args.gitFailUntracked == 'no' else 'normal'
        uncommited = subprocess.check_output("git status --porcelain --untracked=" + untracked, shell=True)
        #print "uncommited: ", uncommited
        if gitclean and uncommited:
            raise NoCleanStateException("new: uncommited files -- please commit changes first\n" + uncommited)
        expConfig['gitHash'], expConfig['gitDescription'] = get_commit()
        if not gitclean and uncommited: expConfig['gitHash'] += '-sloppy'
        expConfig['timestamp'] = datetime.datetime.now().isoformat().split('.')[0]
        expConfig['user'] = getpass.getuser()
        expConfig['shortDescription'] = args.description if args.description else ""
        expConfig['longDescription'] = "" # edit into .mrl file directly
        # make dir and symlink  and
        expDir = self._getExpDir(expId, True)
        if self.lastExpId: os.remove(join(self.outdir,'last'))
        os.symlink(relpath(expDir,self.outdir), join(self.outdir, 'last'))
        os.mkdir(expDir)
        # After this point: expDir is made, no more exception throwing! copy config files
        try:
            if args.copyConfigFrom == 'no':
                pass
            else:
                srcDir = self._getExpDir(self._resolveExpId(args.copyConfigFrom))
                self._copyConfigFrom(srcDir, expDir)
        except (InvalidExpIdException, IOError) as e:
            print("Can't copy config files. ")
            print(e)
            print("Still succesfully created new experiment directory.")
        self._saveExpDotmrl(expDir, expConfig)
        self._renderLaunchScript(expDir, expConfig, self._fmtLaunchScriptFile(expId))
        self.lastExpId = expId
        return expDir

    def info(self, args):
        """ load info from experiment id and print it """
        expId = self._resolveExpId(args.expId)
        expConfig = self._getExpConf(expId)
        return "\n".join("{} : {}".format(k,v) for k,v in expConfig.iteritems())

    def ls(self, args):
        for expId in self.expList[::-1]:
            expConfig = self._getExpConf(expId)
            row = self._fmtSingleExp(expId)
            if args.tm:
                row += "\t" + expConfig['timestamp']
            if args.ghash:
                row += "\t" + expConfig['gitHash']
            if args.gdesc:
                row += "\t" + expConfig['gitDescription']
            if args.desc:
                row += "\t" + expConfig['shortDescription']
            print row
        return ""

    def batch(self, args):
        # resolve id
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        expConfig = self._getExpConf(expId)
        # check if already expanded in batch and cancel
        oldBatchList = []
        if 'batchlist' in expConfig:
            if args.replace:
                oldBatchList = expConfig['batchlist']
            else:
                raise BatchException("Experiment {} is already expanded into a batch of length {}".\
                        format(expId, len(expConfig['batchlist'])))
        # make BatchParser object
        with open(join(expDir, cfg.batchTemplFile)) as fh:
            bp = BatchParser(fh.readlines())
        # check if BatchParser output is non-empty
        if not bp.output:
            err = "BatchParser output is empty, are you sure {} is a batch template?"
            err.format(join(expDir, cfg.batchTemplFile))
            raise BatchException(err)
        # generate output directories, write the output config files
        expConfig['batchlist'] = []
        for i, params, fileContent in bp.output:
            expConfig['batchlist'].append(params)
            subExpDir = join(expDir, self._fmtBatchExp(expId=expId, subExpId=i))
            if i > len(oldBatchList):
                os.mkdir(subExpDir)
            self._copyConfigFrom(expDir, subExpDir)
            with open(join(subExpDir, cfg.batchTemplFile), "w") as fh:
                fh.write(fileContent)
                fh.write("\n")
            self._renderLaunchScript(subExpDir, expConfig, self._fmtLaunchScriptFile(expId, i))

        # update the current .mrl file
        self._saveExpDotmrl(expDir, expConfig)
        self._writeqsubFile(expId, len(bp.output))
        os.remove(join(expDir, self._fmtLaunchScriptFile(expId))) # this conf is a template so doesnt make sense to run it
        return "Succesfully generated config files for expId {}.\n{}".format(expId,expConfig['batchlist'])

    def hpcSubmit(self, args):
        # TODO skip the qsub file and qsub all from mrl. Use _getRunLocations function
        # TODO Use custom() for this? Add pwd function
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        remoteExpDir = join(cfg.hpcBasedir, cfg.outdir, self._fmtSingleExp(expId))
        expConfig = self._getExpConf(expId)
        if 'hpcSubmit' in expConfig and not args.replace:
            raise HpcException("hpcSubmit key already in {} .mrl file. Aborting to avoid data loss.".format(expId))
        # scp to copy to remote
        try:
            self._hpc("mkdir {}".format(remoteExpDir))
        except Exception as e:
            if not args.replace:
                raise HpcException("hpc submission - error in mkdir returned nonzero exit. Dir already exists?")
        self._hpc("scp -r {} {}:{}".format(join(expDir,'*'), cfg.hpcServer, remoteExpDir), ssh=False)
        # subprocess ssh then qsub then print qstat
        if 'batchlist' in expConfig:
            hpcRet = self._hpc('cd {}; ./{}'.format(remoteExpDir, cfg.qsubFile))
            print hpcRet.strip().split("\n")
            # every even-indexed line is an id
            try:
                qsubids = [int(x) for i,x in enumerate(hpcRet.strip().split("\n")) if i%2==0]
            except Exception as e: #failure on qsub? save output
                print "Couldn't save qsubids correctly: ", e
                qsubids = hpcRet
        else:
            hpcRet = self._hpc('cd {}; qsub {}'.format(remoteExpDir, self._fmtLaunchScriptFile(expId)))
            print hpcRet
            qsubids = int(hpcRet)
        expConfig['hpcSubmit'] = qsubids
        self._saveExpDotmrl(expDir, expConfig)

    def hpcFetch(self, args):
        expId = self._resolveExpId(args.expId)
        expDir = self._getExpDir(expId)
        remoteExpDir = join(cfg.hpcBasedir, cfg.outdir, self._fmtSingleExp(expId))
        expConfig = self._getExpConf(expId)
        if 'hpcSubmit' not in expConfig:
            raise HpcException("hpcSubmit key not in in {} .mrl file. Was this launched as hpc job?".format(expId))
        # Locations to fetch from
        runLocs = self._getRunLocations(expId, args.subExpId, expConfig, relativeTo=expDir)
        # Exclude datafiles ifneeded
        if not args.data:
            excludes = " ".join(['--exclude={}'.format(datafile) for datafile in cfg.hpcOutputData])
        else:
            excludes = ""
        for loc in runLocs:
            cmd = "rsync -aP {} {}:{}/ {}/".format(excludes, cfg.hpcServer, join(remoteExpDir,loc), join(expDir,loc))
            self._hpc(cmd, ssh=False)

    def custom(self, args):
        customName = args.customName
        cmdTemplates = self.customTemplates[customName]
        expId = self._resolveExpId(args.expId)
        expConfig = self._getExpConf(expId)
        runLocs = self._getRunLocations(expId, args.subExpId, expConfig, relativeTo=self.outdir)
        originalCwd = os.getcwd() # self.basedir if ran through ./mrl
        cmdParams = {'mrlOutdir': self.outdir, 'mrlBasedir': self.basedir}
        cmdParams.update(expConfig)
        cmdParams.update({k:v for k,v in vars(args).items() if v}) # the optional params if supplied
        # Render & run the commands for all experiments
        for relloc in runLocs:
            absloc = join(self.outdir, relloc)
            print "{} for location {}".format(customName, relloc)
            # Make location cmdParams
            cmdParams.update({'relloc': relloc, 'absloc': absloc})
            # Execute commands in their environment
            for cmdEnvTemplate in cmdTemplates:
                cmdEnv = {k:tmpl.render(cmdParams) for k,tmpl in cmdEnvTemplate.items()}
                if 'cwd' in cmdEnv:
                    os.chdir(os.path.expanduser(cmdEnv['cwd']))
                print cmdEnv['command']
                stdout = open(join(absloc, cmdEnv['output']),"w") if 'output' in cmdEnv else None
                subprocess.call(cmdEnv['command'], stdout=stdout, shell=True)
                if 'cwd' in cmdEnv:
                    os.chdir(originalCwd)

    def _hpc(self, cmd, ssh=True):
        if not self.hpcPass:
            self.hpcPass= getpass.getpass("Password for hpc: ")
        if ssh:
            sshcmd = "ssh {}".format(cfg.hpcServer)
            cmd = '"{}"'.format(cmd.replace('"', '\\"'))
        else:
            sshcmd = ""
        mcmd = "sshpass -p '{}' {} {}".format(self.hpcPass, sshcmd, cmd)
        cleancmd = "sshpass -p '{}' {} {}".format('***', sshcmd, cmd)
        if DEBUG: print cleancmd
        try:
            out = subprocess.check_output(mcmd, shell=True)
        except subprocess.CalledProcessError as e:
            # remove pwd from cmd
            err = str(e).replace(mcmd, cleancmd)
            raise HpcException(err)
        return out

    def _getRunLocations(self, expId, subExpId, expConfig, relativeTo=''):
        if 'batchlist' in expConfig:
            if subExpId == 'all':
                locs = [self._fmtBatchExp(expId, i) for i in range(1,len(expConfig['batchlist'])+1)]
            elif subExpId.isdigit():
                if int(subExpId) > len(expConfig['batchlist']):
                    raise SubExpIdException("subExpId {} out of range (batch size {} in expConfig)".format(subExpId,len(expConfig['batchlist'])))
                locs = [self._fmtBatchExp(expId, int(subExpId)), ]
            else:
                # TODO list of subexpids can be handy
                raise SubExpIdException("Don't understand subExpId {}.".format(subExpId))
        else:
            locs = ['']
        expDir = self._getExpDir(expId)
        if relativeTo != expDir:
            locs = [join(expDir, loc) for loc in locs]
            if relativeTo:
                locs = [relpath(loc, relativeTo) for loc in locs]
        return locs

    def _writeqsubFile(self, expId, N):
        # TODO make git clone  and checkout part of this to avoid 100 clones to the subdirs.
        expDir = self._getExpDir(expId)
        command = 'cd {} && qsub {} && cd .. && sleep 0.01 && echo "qsubbed {}"'
        subdirs = [self._fmtBatchExp(expId=expId, subExpId=i) for i in range(1,N+1)]
        lsfiles = [self._fmtLaunchScriptFile(expId, i) for i in range(1,N+1)]
        cmds = [command.format(subdir, lsfile, subdir) for lsfile,subdir in zip(lsfiles,subdirs)]
        cmds.append("")
        open(join(expDir, cfg.qsubFile),"w").write("\n".join(cmds))
        os.chmod(join(expDir, cfg.qsubFile), 0770)

    def _renderLaunchScript(self, expDir, expConfig, fn):
        """ Launch script is jinja template in basedir config.
        It has access to expDir, the global configuration,
        and all experiment specific variables
        in expConfig."""
        # TODO make git clone and checkout part optional here
        jtemp = Template(cfg.launchScriptTempl)
        tparams = {k:getattr(cfg, k) for k in dir(cfg) if '_' not in k}
        tparams.update(expConfig)
{'fraction_non_alphanumeric': 0.11144374981655954, 'fraction_numerical': 0.0027002436089342842, 'mean_word_length': 2.417795165011536, 'pattern_counts': {'":': 2, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '199903', 'n_tokens_mistral': 10289, 'n_tokens_neox': 10051, 'n_words': 2786}