gen2-RC2-generator

#!/usr/bin/env python3

import yaml
import types
import os
import argparse


class gen2generator:

    """ A class to generate LSST Gen2 scipts"""

    def __init__(self, **keys):

        # Load the configurarion
        self.config = types.SimpleNamespace(**keys)

        # Read in yaml dataset conf manifest
        self.load_dataset()

        # Read in yaml templates conf manifest
        self.load_templates()

        # Define rerun
        self.get_rerun()

        # And the outpath
        self.filepath = self.config.filepath.format(week=self.config.week, DMticket=self.config.DMticket)
        if not os.path.exists(self.filepath):
            os.makedirs(self.filepath)
            print(f"# Created dirname:{self.filepath}")

    def load_dataset(self):
        # Load up the manifest
        with open(self.config.dataset) as f:
            self.data = yaml.load(f, Loader=yaml.FullLoader)
        print(f"# Loaded dataset definitions from: {self.config.dataset}")

    def load_templates(self):
        # Load up the manifest
        with open(self.config.templates) as f:
            self.templates = yaml.load(f, Loader=yaml.FullLoader)
        print(f"# Loaded dataset templates from: {self.config.templates}")

    def get_rerun(self):
        # Define the rerun
        self.rerun = self.config.rerun_format.format(week=self.config.week, DMticket=self.config.DMticket)
        print(f"# Defined rerun as: {self.rerun}")

    def get_visit_str(self, tract, filter):
        # get the visit in the format we need
        visit_str = "^".join([str(i) for i in self.data['visits'][tract][filter]])
        return visit_str

    def get_filter_str(self, tract):
        # get the filter in the format we need
        filter_str = "^".join([str(i) for i in self.data['visits'][tract]])
        return filter_str

    def task_filename_tract_filter(self, task, tract, filter, ext=''):
        index = self.templates['tasks_to_execute_ordered'].index(task) + 1
        name = os.path.join(self.filepath, f"{index:02d}_{task}_{tract}_{filter}{ext}")
        return name

    def task_filename_tract(self, task, tract, ext=''):
        index = self.templates['tasks_to_execute_ordered'].index(task) + 1
        name = os.path.join(self.filepath, f"{index:02d}_{task}_{tract}{ext}")
        return name

    def task_filename_task(self, task, ext=''):
        index = self.templates['tasks_to_execute_ordered'].index(task) + 1
        name = os.path.join(self.filepath, f"{index:02d}_{task}{ext}")
        return name

    def task_per_task(self, task, ext=''):

        """ Format task per rerun/task only """
        outfile = self.task_filename_task(task, ext=ext)
        out = open(outfile, "w")
        out.write(self.templates['templates'][task].format(rerun=self.rerun))
        out.close()
        print(f"# file: {outfile} is ready")

    def task_per_tract(self, task, ext=''):

        """ Format a task per tract only """

        for tract in self.data['visits'].keys():

            outfile = self.task_filename_tract(task, tract, ext=ext)
            out = open(outfile, "w")
            filter_str = self.get_filter_str(tract)
            kw = {'rerun': self.rerun, 'filter_str': filter_str, 'tract': tract}
            if task in self.templates['tasks_specs']:
                kw.update(self.templates['tasks_specs'][task][tract])
            out.write(self.templates['templates'][task].format(**kw))
            out.close()
            print(f"# file: {outfile} is ready")

    def task_per_tract_filter(self, task, ext=''):

        """ Format a task per tract and filter"""

        outfile = None
        head_task = False
        head_tract = False
        kw = {'rerun': self.rerun, 'rerun2': self.config.rerun2}

        # Head per rerun
        if 'head_task' in self.templates['templates'][task]:
            outfile = self.task_filename_task(task, ext=ext)
            out = open(outfile, "w")
            out.write(self.templates['templates'][task]['head_task'].format(**kw))
            head_task = True

        for tract in self.data['visits'].keys():

            # Head per tract
            if 'head_tract' in self.templates['templates'][task]:
                outfile = self.task_filename_tract(task, tract, ext=ext)
                out = open(outfile, "w")
                kw.update({'tract': tract})
                out.write(self.templates['templates'][task]['head_tract'].format(**kw))
                head_tract = True

            for filter in self.data['visits'][tract].keys():

                if not head_task and not head_tract:
                    outfile = self.task_filename_tract(task, tract, ext=ext)
                    out = open(outfile, "w")

                # get the visit in the format we need
                visit = self.get_visit_str(tract, filter)
                kw.update({'visit': visit, 'filter': filter, 'tract': tract})
                if task in self.templates['tasks_specs']:
                    kw.update(self.templates['tasks_specs'][task][tract])

                if 'multi' in self.templates['templates'][task]:
                    out.write(self.templates['templates'][task]['multi'].format(**kw) + '\n')
                else:
                    out.write(self.templates['templates'][task].format(**kw) + '\n')

                # close file per filter
                if not head_task and not head_tract:
                    out.close()
                    print(f"# file: {outfile} is ready")

            # Out the loop print
            if 'no_loop' in self.templates['templates'][task]:
                out.write(self.templates['templates'][task]['no_loop'].format(**kw))

            # close file accordingly by tract
            if head_tract:
                out.close()
                print(f"# file: {outfile} is ready")

        # close file accordingly by task
        if head_task:
            out.close()
            print(f"# file: {outfile} is ready")

    def task_per_tract_filter_visit_multi(self, task, ext=''):

        """ Format a task per tract/filter/visit for multi-progran slurm"""

        kw = {'rerun': self.rerun, 'rerun2': self.config.rerun2}

        # Loop over tracts
        for tract in self.data['visits'].keys():

            # Loop over filters
            for filter in self.data['visits'][tract].keys():

                # One file per filter and tract
                outfile = self.task_filename_tract_filter(task, tract, filter, ext=ext)
                out = open(outfile, "w")
                # get the visit in the format we need
                tractname = self.data['tractname'][tract]
                visit = self.get_visit_str(tract, filter)
                ntasks = len(self.data['visits'][tract][filter])
                mpfile = f"{task}_{tract}_{tractname}_{filter}.mp"

                # Create the dict with kwargs
                kw.update({'visit': visit, 'filter': filter, 'tract': tract,
                           'ntasks': ntasks, 'tractname': tractname, 'mpfile': mpfile})
                # Extra info per task
                if task in self.templates['tasks_specs']:
                    kw.update(self.templates['tasks_specs'][task][tract])

                # The head subtemplate of the slurm
                out.write(self.templates['templates'][task]['head_filter'].format(**kw))
                out.close()
                print(f"# file: {outfile} is ready")

                # Loop over each visit
                # The multiprog template of the slurm goes on separate file
                outfile = self.task_filename_tract_filter(task, tract, filter, ext='.mp')
                out = open(outfile, "w")
                mpindex = 0
                for visitID in self.data['visits'][tract][filter]:
                    kw.update({'mpindex': mpindex, 'visitID': visitID})
                    out.write(self.templates['templates'][task]['multi'].format(**kw) + '\n')
                    mpindex += 1
                out.close()
                print(f"# file: {outfile} is ready")

    def do_tasks(self):

        # Loop over tasks to execute
        for task in self.templates['tasks_to_execute_ordered']:
            ext = self.templates['tasks_extname'][task]
            print(f"# --- Preparing task: {task} ---- ")
            if task in self.templates['tasks_to_execute_slurm_multi']:
                self.task_per_tract_filter_visit_multi(task, ext=ext)
            elif task in self.templates['tasks_per_tract_filter']:
                self.task_per_tract_filter(task, ext=ext)
            elif task in self.templates['tasks_per_tract']:
                self.task_per_tract(task, ext=ext)
            elif task in self.templates['tasks_per_task']:
                self.task_per_task(task, ext=ext)
            else:
                exit(f"ERROR: task: {task} NOT defined on task type")


def cmdline():

    parser = argparse.ArgumentParser(description="Gen2 Generation script")
    parser.add_argument("--DMticket", action="store", required=True,
                        help="The DM-ticket number [i.e.: DM-XXXX]")
    parser.add_argument("--week", action="store", required=True,
                        help="The build week [i.e. w_2019_38]")
    parser.add_argument("--rerun2", action="store", required=True,
                        help="Name of the rerun to compare with")
    parser.add_argument("--dataset", action="store", default='config/RC2dataset.yaml',
                        help="yaml file with dataset description")
    parser.add_argument("--templates", action="store", default='config/RC2templates.yaml',
                        help="yaml file with template description")
    parser.add_argument("--rerun_format", action="store", default="RC/{week}/{DMticket}",
                        help="The format for the rerun path")
    parser.add_argument("--filepath", action="store", default="rerun_scripts/{DMticket}",
                        help="The format for the output script files")
    args = parser.parse_args()
    return args


if __name__ == "__main__":

    args = cmdline()
    gen2 = gen2generator(**args.__dict__)
    gen2.do_tasks()