pash_compiler.py

import argparse
import sys
import pickle
import traceback
from datetime import datetime

from pash_annotations.annotation_generation.datatypes.parallelizability.AggregatorKind import AggregatorKindEnum

import config
import env_vars_util
from ir import *
from ast_to_ir import compile_asts
from ir_to_ast import to_shell
from pash_graphviz import maybe_generate_graphviz
from util import *

from definitions.ir.aggregator_node import *

from definitions.ir.dfg_node import DFGNode
from definitions.ir.nodes.eager import *
from definitions.ir.nodes.pash_split import *

import definitions.ir.nodes.r_merge as r_merge
import definitions.ir.nodes.r_split as r_split
import definitions.ir.nodes.r_unwrap as r_unwrap
import definitions.ir.nodes.dgsh_tee as dgsh_tee
import definitions.ir.nodes.dfs_split_reader as dfs_split_reader
# Distirbuted Exec
import dspash.hdfs_utils as hdfs_utils 
# Serverless Exec
# import serverless.serverless_executor as splash

runtime_config = {}
## We want to catch all exceptions here so that they are logged correctly
## and not just printed to the stderr.
def main():
    try:
        main_body()
    except Exception:
        log("Compiler failed, no need to worry, executing original script...")
        log(traceback.format_exc())
        sys.exit(1)

def main_body():
    global runtime_config

    ## Parse arguments
    args = parse_args()
    config.set_config_globals_from_pash_args(args)

    ## Load the configuration
    if not config.config:
        config.load_config(args.config_path)

    runtime_config = config.config['distr_planner']

    ## Read any shell variables files if present
    vars_dict = env_vars_util.read_vars_file(args.var_file)
    config.set_vars_file(args.var_file, vars_dict)

    log("Input:", args.input_ir, "Compiled file:", args.compiled_script_file)

    ## Call the main procedure
    compiler_config = CompilerConfig(args.width)
    ast_or_ir = compile_optimize_output_script(args.input_ir, args.compiled_script_file, args, compiler_config)
    maybe_generate_graphviz(ast_or_ir, args)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("compiled_script_file",
                        help="the file in which to output the compiled script")
    parser.add_argument("input_ir",
                        help="the file containing the dataflow graph to be optimized and executed")
    parser.add_argument("--var_file",
                        help="determines the path of a file containing all shell variables.",
                        default=None)
    config.add_common_arguments(parser)
    args, unknown_args = parser.parse_known_args()
    return args

## TODO: Add more fields from args in this
class CompilerConfig:
    def __init__(self, width):
        self.width = width
    
    def __repr__(self):
        return f'CompilerConfig(Width:{self.width})'

def compile_ir(ir_filename, compiled_script_file, args, compiler_config):
    """
    Return IR object for compilation success. None otherwise.
    """
    ret = None
    try:
        ret = compile_optimize_output_script(ir_filename, compiled_script_file, args, compiler_config)
    except Exception as e:
        log("WARNING: Exception caught:", e)
        # traceback.print_exc()

    return ret

def compile_optimize_output_script(ir_filename, compiled_script_file, args, compiler_config):
    global runtime_config
    
    ret = None

    ## Load the df_region from a file
    candidate_df_region = load_df_region(ir_filename)
    
    ## Compile it
    optimized_ast_or_ir = compile_optimize_df_region(candidate_df_region, args, compiler_config)

    ## Call the backend that executes the optimized dataflow graph
    ## TODO: Should never be the case for now. This is obsolete.
    assert(not runtime_config['distr_backend'])

    ## If the candidate DF region was indeed a DF region then we have an IR
    ## which should be translated to a parallel script.
    if(isinstance(optimized_ast_or_ir, IR)):
        if args.distributed_exec:
            ir_filename = ptempfile()
            script_to_execute = f"$PASH_TOP/compiler/dspash/remote_exec_graph.sh {ir_filename}\n"
            ## This might not be needed anymore (since the output script is output anyway)
            ## TODO: This is probably useless, remove
            maybe_log_optimized_script(script_to_execute, args)

            with open(ir_filename, "wb") as f:
                obj = (optimized_ast_or_ir, config.config['shell_variables'])
                pickle.dump(obj, f)
        elif args.serverless_exec:
            ir_filename = ptempfile()
            script_to_execute = f"$PASH_TOP/compiler/serverless/serverless_exec_graph.sh {ir_filename}\n"
            with open(ir_filename, "wb") as f:
                obj = (optimized_ast_or_ir, config.config['shell_variables'], args)
                pickle.dump(obj, f)
        else:
            script_to_execute = to_shell(optimized_ast_or_ir, args)
            
        log("Optimized script saved in:", compiled_script_file)
        with open(compiled_script_file, "w") as f:
            f.write(script_to_execute)
    
        ret = optimized_ast_or_ir
    else:
        raise Exception("Script failed to compile!")
    
    return ret

def load_df_region(ir_filename):
    log("Retrieving candidate DF region: {} ... ".format(ir_filename), end='')
    with open(ir_filename, "rb") as ir_file:
        candidate_df_region = pickle.load(ir_file)
    log("Done!")
    return candidate_df_region

def compile_optimize_df_region(df_region, args, compiler_config):
    ## Compile the candidate DF regions
    compilation_start_time = datetime.now()
    asts_and_irs = compile_candidate_df_region(df_region, config.config)
    compilation_end_time = datetime.now()
    print_time_delta("Compilation", compilation_start_time, compilation_end_time)

    ## Optimize all the IRs that can be optimized
    if(args.no_optimize):
        optimized_asts_and_irs = asts_and_irs
    else:
        optimized_asts_and_irs = optimize_irs(asts_and_irs, args, compiler_config)

    ## TODO: Normally this could return more than one compiled ASTs (containing IRs in them).
    ##       To correctly handle that we would need to really replace the optimized IRs
    ##       with the final parallel corresponding scripts.
    ##
    ##       However, for now we just assume that there is one IR that we can execute as is.
    ##
    ## TODO: This might bite us with the quick-abort.
    ##       It might complicate things having a script whose half is compiled to a graph and its other half not.
    assert(len(optimized_asts_and_irs) == 1)
    optimized_ast_or_ir = optimized_asts_and_irs[0]
    
    return optimized_ast_or_ir

def maybe_log_optimized_script(script_to_execute, args):
    ## TODO: Merge this write with the one below. Maybe even move this logic in `pash_runtime.sh`
    ## Output the optimized shell script for inspection
    if(args.output_optimized):
        output_script_path = runtime_config['optimized_script_filename']
        with open(output_script_path, "w") as output_script_file:
            log("Optimized script:")
            log(script_to_execute)
            output_script_file.write(script_to_execute)

def compile_candidate_df_region(candidate_df_region, config):
    ## This is for the files in the IR
    fileIdGen = FileIdGen()
    
    ## If the candidate DF region is not from the top level then
    ## it won't be a list and thus we need to make it into a list to compile it.
    if(not isinstance(candidate_df_region, list)):
        candidate_df_region = [candidate_df_region]

    ## Compile the asts
    ## TODO: Since compilation happens at runtime, we can now expand everything accordingly.
    ##       We can do that using a shell for start:
    ##         if a word is safe to expand, then call a shell to expand it.
    compiled_asts = compile_asts(candidate_df_region, fileIdGen, config)

    return compiled_asts

## TODO: Switch args to compiler_config
def optimize_irs(asts_and_irs, args, compiler_config):
    global runtime_config

    optimization_start_time = datetime.now()

    optimized_asts_and_irs = []
    for ast_or_ir in asts_and_irs:
        if(isinstance(ast_or_ir, IR)):
            ## Assert that the graph that was returned from compilation is valid
            assert(ast_or_ir.valid())

            # log(ir_node)
            # with cProfile.Profile() as pr:
            distributed_graph = choose_and_apply_parallelizing_transformations(ast_or_ir, compiler_config.width,
                                                                      runtime_config['batch_size'],
                                                                      args.r_split_batch_size)
            # pr.print_stats()

            # Eagers are added in remote notes when using distributed exec
            if(not args.no_eager and not args.distributed_exec): 
                eager_distributed_graph = add_eager_nodes(distributed_graph)
            else:
                eager_distributed_graph = distributed_graph

            ## Assert that the graph stayed valid after all transformations
            assert(eager_distributed_graph.valid())

            ## Print statistics of output nodes
            print_graph_statistics(eager_distributed_graph)

            optimized_asts_and_irs.append(eager_distributed_graph)
        else:
            optimized_asts_and_irs.append(ast_or_ir)

    optimization_end_time = datetime.now()
    print_time_delta("Optimization", optimization_start_time, optimization_end_time)

    return optimized_asts_and_irs


def print_graph_statistics(graph):
    total_nodes = graph.nodes
    eager_nodes = [node for node in total_nodes.values() if isinstance(node, Eager)]
    log("Total nodes after optimization:", len(total_nodes))
    log(" -- out of which:")
    log("Eager nodes:", len(eager_nodes))


def choose_and_apply_parallelizing_transformations(graph, fan_out, batch_size, r_split_batch_size):
    parallelizer_map = choose_parallelizing_transformations(graph)
    apply_parallelizing_transformations(graph, parallelizer_map, fan_out, batch_size, 
                                        r_split_batch_size)
    return graph


def choose_parallelizing_transformations(graph): # shall return map
    source_node_ids = graph.source_nodes()
    parallelizer_map = {}
    workset = source_node_ids
    visited = set()
    # We apply a modified BFS such that we ensure that we know which parallelizer was chosen for all previous nodes
    # and assume that the decision for any subsequent node will exploit any potential synergy effects
    while (len(workset) > 0):
        curr_id = workset.pop(0)
        assert(isinstance(curr_id, int))
        all_previous_nodes_visited = all(prev in visited for prev in graph.get_previous_nodes(curr_id))
        if not all_previous_nodes_visited:
            workset.append(curr_id)
        elif not curr_id in visited:
            next_node_ids = graph.get_next_nodes(curr_id)
            workset += next_node_ids
            parallelizer_map[curr_id] = choose_parallelizing_transformation(curr_id, graph)
            visited.add(curr_id)
    return parallelizer_map


## This currently chooses the best parallelization based on priority:
## 1. The round robin
## 2. The round robin after having performed unwrap (not sure why this is the second priority)
## 3. The consecutive chunks
## 
## TODO: In the future, we could develop more complex strategies      
def choose_parallelizing_transformation(curr_id, graph): # shall return map entry
    curr = graph.get_node(curr_id)
    list_all_parallelizers_in_priority = [curr.get_option_implemented_round_robin_parallelizer(),
                                          curr.get_option_implemented_round_robin_with_unwrap_parallelizer(),
                                          curr.get_option_implemented_consecutive_chunks_parallelizer()]
    return next((item for item in list_all_parallelizers_in_priority if item is not None), None)


def apply_parallelizing_transformations(graph, parallelizer_map, fan_out, batch_size, r_split_batch_size):
    fileIdGen = graph.get_file_id_gen()
    node_id_non_none_parallelizer_list = [(node_id, parallelizer) for (node_id, parallelizer) in parallelizer_map.items()
                                                                  if parallelizer is not None]
    for (node_id, parallelizer) in node_id_non_none_parallelizer_list:
        graph.apply_parallelization_to_node(node_id, parallelizer, fileIdGen, fan_out, r_split_batch_size)

def split_hdfs_cat_input(hdfs_cat, next_node, graph, fileIdGen):
    """
    Replaces hdfs cat with a cat per block, each cat uses has an HDFSResource input fid
    Returns: A normal Cat that merges the blocks (will be removed when parallizing next_node)
    """
    assert(isinstance(hdfs_cat, HDFSCat))

    ## At the moment this only works for nodes that have one standard input.
    if len(next_node.get_standard_inputs()) != 1:
        return

    hdfscat_input_id = hdfs_cat.get_standard_inputs()[0]
    hdfs_fid = graph.get_edge_fid(hdfscat_input_id)
    hdfs_filepath = str(hdfs_fid.get_resource())
    output_ids = []

    # Create a cat command per file block
    file_config = hdfs_utils.get_file_config(hdfs_filepath)
    dummy_config_path = ptempfile() # Dummy config file, should be updated by workers
    for split_num, block in enumerate(file_config.blocks):
        resource = DFSSplitResource(file_config.dumps(), dummy_config_path, split_num, block.hosts)
        block_fid = fileIdGen.next_file_id()
        block_fid.set_resource(resource)
        graph.add_edge(block_fid)

        output_fid = fileIdGen.next_file_id()
        output_fid.make_ephemeral()
        output_ids.append(output_fid.get_ident())
        graph.add_edge(output_fid)

        split_reader_node = dfs_split_reader.make_dfs_split_reader_node([block_fid.get_ident()], output_fid.get_ident(), split_num, config.HDFS_PREFIX)
        graph.add_node(split_reader_node)

    # Remove the HDFS Cat command as it's not used anymore
    graph.remove_node(hdfs_cat.get_id())

    ## input of next command is output of new merger.
    input_id = next_node.get_standard_inputs()[0]
    new_merger = make_cat_node(output_ids, input_id)
    graph.add_node(new_merger)

    return new_merger



## This functions adds an eager on a given edge.
def add_eager(eager_input_id, graph, fileIdGen):
    new_fid = fileIdGen.next_ephemeral_file_id()
    new_id = new_fid.get_ident()

    ## TODO: seperate to better use dgsh-tee params and maybe deprecate eager
    eager_node = dgsh_tee.make_dgsh_tee_node(eager_input_id, new_id)

    ## Add the edges and the nodes to the graph
    graph.add_edge(new_fid)

    ## Modify the next node inputs to be the new inputs
    next_node_id = graph.edges[eager_input_id][2]
    if(not next_node_id is None):
        next_node = graph.get_node(next_node_id)
        next_node.replace_edge(eager_input_id, new_id)
        graph.set_edge_to(new_id, next_node_id)

    graph.add_node(eager_node)


## This function adds eager nodes wherever the width of graph is
## becoming smaller.
def add_eager_nodes(graph):
    source_node_ids = graph.source_nodes()

    ## Generate a fileIdGen that doesnt clash with graph fids.
    fileIdGen = graph.get_file_id_gen()

    ## Get the next nodes
    workset = [node for source_node_id in source_node_ids for node in graph.get_next_nodes(source_node_id)]
    visited = set()
    while (len(workset) > 0):
        curr_id = workset.pop(0)
        curr = graph.get_node(curr_id)
        if (not curr_id in visited):
            visited.add(curr_id)
            next_node_ids = graph.get_next_nodes(curr_id)
            workset += next_node_ids

            ## TODO: Make sure that we don't add duplicate eager nodes

            ## Add eager nodes if the node has more than one input
            curr_input_ids = graph.get_node_input_ids(curr_id)
            if (len(curr_input_ids) > 1):
                ## TODO: If we know that a command reads its inputs in a list,
                ##       then we might not need to put an eager on its first input.
                ## Note: This cannot be done for `sort -m` so we need to know in the
                ##       annotations whether input consumption is in order or not.

                for curr_input_id in curr_input_ids:
                    _fid, from_node, to_node = graph.edges[curr_input_id]
                    assert(to_node == curr_id)
                    ## If the edge is an input edge, then we don't want to put eager.
                    if(not from_node is None):
                        add_eager(curr_input_id, graph, fileIdGen)

            if(isinstance(curr, Split)):
                eager_input_ids = curr.get_output_list()[:-1]
                for edge_id in eager_input_ids:
                    add_eager(edge_id, graph, fileIdGen)

            ## Add an eager after r_unwrap            
            if(isinstance(curr, r_unwrap.RUnwrap)):
                eager_input_id = curr.get_output_list()[0]
                add_eager(eager_input_id, graph, fileIdGen)

            ## Add an eager after r_split
            if(isinstance(curr, r_split.RSplit)):
                eager_input_ids = curr.get_output_list()
                for edge_id in eager_input_ids:
                    add_eager(edge_id, graph, fileIdGen)

    return graph


if __name__ == "__main__":
    main()