diff --git a/esm_runscripts/__init__.py b/esm_runscripts/__init__.py index 3a4c55d..cdbe129 100644 --- a/esm_runscripts/__init__.py +++ b/esm_runscripts/__init__.py @@ -2,7 +2,7 @@ __author__ = """Dirk Barbi""" __email__ = 'dirk.barbi@awi.de' -__version__ = "5.0.14" +__version__ = "5.0.17" from .sim_objects import * from .batch_system import * @@ -11,6 +11,7 @@ from .compute import * from .tidy import * from .prepare import * +from .last_minute import * from .postprocess import * from .filelists import * from .tidy import * diff --git a/esm_runscripts/batch_system.py b/esm_runscripts/batch_system.py index 22ff9b7..dcd6a8d 100644 --- a/esm_runscripts/batch_system.py +++ b/esm_runscripts/batch_system.py @@ -1,4 +1,5 @@ import os +import textwrap import sys import esm_environment @@ -60,20 +61,31 @@ def get_batch_header(config): this_batch_system = config["computer"] if "sh_interpreter" in this_batch_system: header.append("#!" + this_batch_system["sh_interpreter"]) - tasks = batch_system.calculate_requirements(config) + tasks, nodes = batch_system.calculate_requirements(config) replacement_tags = [("@tasks@", tasks)] - all_flags = [ - "partition_flag", - "time_flag", - "tasks_flag", - "output_flags", - "name_flag", - ] + if config["general"].get("taskset", False): + replacement_tags = [("@nodes@", nodes)] + all_flags = [ + "partition_flag", + "time_flag", + "nodes_flag", + "output_flags", + "name_flag", + ] + else: + all_flags = [ + "partition_flag", + "time_flag", + "tasks_flag", + "output_flags", + "name_flag", + ] conditional_flags = [ "accounting_flag", "notification_flag", "hyperthreading_flag", "additional_flags", + "overcommit_flag" ] if config["general"]["jobtype"] in ["compute", "tidy_and_resume"]: conditional_flags.append("exclusive_flag") @@ -93,10 +105,13 @@ def get_batch_header(config): @staticmethod def calculate_requirements(config): tasks = 0 + nodes = 0 if config["general"]["jobtype"] == "compute": for model in config["general"]["valid_model_names"]: if "nproc" in config[model]: tasks += config[model]["nproc"] + if config["general"].get("taskset", False): + nodes +=int((config[model]["nproc"]*config[model]["omp_num_threads"])/config['computer']['cores_per_node']) elif "nproca" in config[model] and "nprocb" in config[model]: tasks += config[model]["nproca"] * config[model]["nprocb"] @@ -111,7 +126,7 @@ def calculate_requirements(config): elif config["general"]["jobtype"] == "post": tasks = 1 - return tasks + return tasks, nodes @staticmethod def get_environment(config): @@ -119,6 +134,39 @@ def get_environment(config): env = esm_environment.environment_infos("runtime", config) return env.commands + @staticmethod + def determine_nodelist(config): + setup_name = config['general']['setup_name'] + if config['general'].get('multi_srun'): + for run_type in config['general']['multi_srun']: + print(run_type) + total_tasks = 0 + for model in config['general']['multi_srun'][run_type]['models']: + print(total_tasks) + # determine how many nodes that component needs + if "nproc" in config[model]: + print("Adding to total_tasks") + total_tasks += int(config[model]["nproc"]) + print(total_tasks) + elif "nproca" in config[model] and "nprocb" in config[model]: + print("Adding to total_tasks") + total_tasks += int(config[model]["nproca"])*int(config[model]["nprocb"]) + print(total_tasks) + + # KH 30.04.20: nprocrad is replaced by more flexible + # partitioning using nprocar and nprocbr + if "nprocar" in config[model] and "nprocbr" in config[model]: + if config[model]["nprocar"] != "remove_from_namelist" and config[model]["nprocbr"] != "remove_from_namelist": + print("Adding to total_tasks") + total_tasks += config[model]["nprocar"] * config[model]["nprocbr"] + print(total_tasks) + + else: + continue + config['general']['multi_srun'][run_type]['total_tasks'] = total_tasks + print(config['general']['multi_srun']) + + @staticmethod def get_extra(config): extras = [] @@ -151,11 +199,16 @@ def get_run_commands(config): # here or in compute.py? commands.append( "echo " + line + " >> " + config["general"]["experiment_log_file"] ) + if config['general'].get('multi_srun'): + return get_run_commands_multisrun(config, commands) commands.append("time " + batch_system["execution_command"] + " &") return commands + + @staticmethod def get_submit_command(config, sadfilename): + # FIXME(PG): Here we need to include a multi-srun thing commands = [] batch_system = config["computer"] if "submit" in batch_system: @@ -175,6 +228,8 @@ def write_simple_runscript(config): sadfilename = batch_system.get_sad_filename(config) header = batch_system.get_batch_header(config) environment = batch_system.get_environment(config) + # NOTE(PG): This next line allows for multi-srun simulations: + batch_system.determine_nodelist(config) extra = batch_system.get_extra(config) if config["general"]["verbose"]: @@ -200,6 +255,11 @@ def write_simple_runscript(config): print("ERROR -- Not sure if you were in a contained or open run!") print("ERROR -- See write_simple_runscript for the code causing this.") sys.exit(1) + + if "modify_config_file_abspath" in config["general"]: + if config["general"]["modify_config_file_abspath"]: + tidy_call += " -m " + config["general"]["modify_config_file_abspath"] + elif config["general"]["jobtype"] == "post": tidy_call = "" commands = config["general"]["post_task_list"] @@ -214,12 +274,40 @@ def write_simple_runscript(config): sadfile.write(line + "\n") sadfile.write("\n") sadfile.write("cd " + config["general"]["thisrun_work_dir"] + "\n") + if config["general"].get("taskset", False): + sadfile.write("\n"+"#Creating hostlist for MPI + MPI&OMP heterogeneous parallel job" + "\n") + sadfile.write("rm -f ./hostlist" + "\n") + sadfile.write(f"export SLURM_HOSTFILE={config['general']['thisrun_work_dir']}/hostlist\n") + sadfile.write("IFS=$'\\n'; set -f" + "\n") + sadfile.write("listnodes=($(< <( scontrol show hostnames $SLURM_JOB_NODELIST )))"+"\n") + sadfile.write("unset IFS; set +f" + "\n") + sadfile.write("rank=0" + "\n") + sadfile.write("current_core=0" + "\n") + sadfile.write("current_core_mpi=0" + "\n") + for model in config["general"]["valid_model_names"]: + if model != "oasis3mct": + sadfile.write("mpi_tasks_"+model+"="+str(config[model]["nproc"])+ "\n") + sadfile.write("omp_threads_"+model+"="+str(config[model]["omp_num_threads"])+ "\n") + import pdb + #pdb.set_trace() + sadfile.write("for model in " + str(config["general"]["valid_model_names"])[1:-1].replace(',', '').replace('\'', '') +" ;do"+ "\n") + sadfile.write(" eval nb_of_cores=\${mpi_tasks_${model}}" + "\n") + sadfile.write(" eval nb_of_cores=$((${nb_of_cores}-1))" + "\n") + sadfile.write(" for nb_proc_mpi in `seq 0 ${nb_of_cores}`; do" + "\n") + sadfile.write(" (( index_host = current_core / " + str(config["computer"]["cores_per_node"]) +" ))" + "\n") + sadfile.write(" host_value=${listnodes[${index_host}]}" + "\n") + sadfile.write(" (( slot = current_core % " + str(config["computer"]["cores_per_node"]) +" ))" + "\n") + sadfile.write(" echo $host_value >> hostlist" + "\n") + sadfile.write(" (( current_core = current_core + omp_threads_${model} ))" + "\n") + sadfile.write(" done" + "\n") + sadfile.write("done" + "\n\n") for line in commands: sadfile.write(line + "\n") sadfile.write("process=$! \n") sadfile.write("cd " + config["general"]["experiment_scripts_dir"] + "\n") sadfile.write(tidy_call + "\n") + config["general"]["submit_command"] = batch_system.get_submit_command( config, sadfilename ) @@ -234,8 +322,29 @@ def write_simple_runscript(config): six.print_("Contents of ", self.bs.filename, ":") with open(self.bs.filename, "r") as fin: print(fin.read()) + + # Write the environment in a file that can be sourced from preprocessing and + # postprocessing scripts + batch_system.write_env(config, environment, sadfilename) + return config + @staticmethod + def write_env(config, environment, sadfilename): + folder = config["general"]["thisrun_scripts_dir"] + this_batch_system = config["computer"] + sadfilename_short = sadfilename.split("/")[-1] + envfilename = folder + "/env.sh" + + with open(envfilename, "w") as envfile: + if "sh_interpreter" in this_batch_system: + envfile.write("#!" + this_batch_system["sh_interpreter"] + "\n") + envfile.write(f"# ENVIRONMENT used in {sadfilename_short}\n") + envfile.write("# Use this file to source the environment in your\n") + envfile.write("# preprocessing or postprocessing scripts\n\n") + for line in environment: + envfile.write(line + "\n") + @staticmethod def submit(config): if not config["general"]["check"]: @@ -256,3 +365,112 @@ def submit(config): ) print() return config + + +def get_run_commands_multisrun(config, commands): + default_exec_command = config['computer']["execution_command"] + print("---> This is a multi-srun job.") + print("The default command:") + print(default_exec_command) + print("Will be replaced") + # Since I am already confused, I need to write comments. + # + # The next part is actually a shell script fragment, which will be injected + # into the "sad" file. sad = Sys Admin Dump. It's sad :-( + # + # In this part, we figure out what compute nodes we are using so we can + # specify nodes for each srun command. That means, ECHAM+FESOM will use one + # pre-defined set of nodes, PISM another, and so on. That should be general + # enough to also work for other model combos... + # + # Not sure if this is specific to Mistral as a HPC, Slurm as a batch + # system, or whatever else might pop up... + # @Dirk, please move this where you see it best (I guess slurm.py) + job_node_extraction = r""" + # Job nodes extraction + nodeslurm=$SLURM_JOB_NODELIST + echo "nodeslurm = ${nodeslurm}" + # Get rid of the hostname and surrounding brackets: + tmp=${nodeslurm#"*["} + nodes=${tmp%]*} + # Turn it into an array seperated by newlines: + myarray=(`echo ${nodes} | sed 's/,/\n/g'`) + # + idx=0 + for element in "${myarray[@]}"; do + if [[ "$element" == *"-"* ]]; then + array=(`echo $element | sed 's/-/\n/g'`) + for node in $(seq ${array[0]} ${array[1]}); do + nodelist[$idx]=${node} + idx=${idx}+1 + done + else + nodelist[$idx]=${element} + idx=${idx}+1 + fi + done + + for element in "${nodelist[@]}"; do + echo "${element}" + done + """ + + def assign_nodes(run_type, need_length=False, start_node=0, num_nodes_first_model=0): + template = f""" + # Assign nodes for {run_type} + {run_type}="" + %%NEED_LENGTH%% + for idx in $srbseq {start_node} $srbsrb???-1erberberb; do + if ssbssb $idx == $srbsrb???-1erberb esbesb; then + {run_type}="$scb{run_type}ecb$scbnodelist[$idx]ecb" + else + {run_type}="$scb{run_type}ecb$scbnodelistssb$idxesbecb," + fi + done + echo "{run_type} nodes: $scb{run_type}ecb" + """ + # Since Python f-strings and other braces don't play nicely together, + # we replace some stuff: + # + # For the confused: + # scb = start curly brace { + # ecb = end curly brace } + # ssb = start square brace [ + # esb = end square brace ] + # srb = start round brace ( + # erb = end round brace ) + template = template.replace("scb", "{") + template = template.replace("ecb", "}") + template = template.replace("ssb", "[") + template = template.replace("esb", "]") + template = template.replace("srb", "(") + template = template.replace("erb", ")") + # Get rid of the starting spaces (they come from Python as the string + # is defined inside of this function which is indented (facepalm)) + template = textwrap.dedent(template) + # TODO: Some replacements + if need_length: + length_stuff = r"length=${#nodelist[@]}" + template = template.replace("%%NEED_LENGTH%%", length_stuff) + template = template.replace("???", "length") + else: + template = template.replace("%%NEED_LENGTH%%", "") + template = template.replace("???", str(num_nodes_first_model)) + return template + + + commands.append(textwrap.dedent(job_node_extraction)) + for idx, run_type in enumerate(config['general']['multi_srun']): + if idx == 0: + start_node = run_type + num_nodes_first_model = config['general']['multi_srun'][run_type]['total_tasks'] / config['computer']['cores_per_node'] + num_nodes_first_model = int(num_nodes_first_model) + nodes = assign_nodes(run_type, need_length=False, num_nodes_first_model=num_nodes_first_model) + else: + nodes = assign_nodes(run_type, need_length=True, start_node=start_node) + commands.append(nodes) + for run_type in config['general']['multi_srun']: + new_exec_command = default_exec_command.replace("hostfile_srun", config['general']['multi_srun'][run_type]['hostfile']) + new_exec_command += f" --nodelist ${run_type}" + commands.append("time " + new_exec_command + " &") + return commands diff --git a/esm_runscripts/cli.py b/esm_runscripts/cli.py index d95ec7f..0736a0d 100644 --- a/esm_runscripts/cli.py +++ b/esm_runscripts/cli.py @@ -62,6 +62,14 @@ def parse_shargs(): action="store_true", ) + parser.add_argument( + "--modify-config", + "-m", + dest="modify", + help="[m]odify configuration", + default="", # kh 15.07.20 "usermods.yaml" + ) + parser.add_argument( "-j", "--last_jobtype", @@ -122,6 +130,7 @@ def main(): verbose = False inspect = None use_venv = None + modify_config_file = None parsed_args = vars(ARGS) @@ -153,10 +162,8 @@ def main(): use_venv = parsed_args["contained_run"] if parsed_args["open_run"] is not None: use_venv = not parsed_args["open_run"] - - - - + if "modify" in parsed_args: + modify_config_file = parsed_args["modify"] command_line_config = {} command_line_config["check"] = check @@ -170,6 +177,8 @@ def main(): command_line_config["verbose"] = verbose command_line_config["inspect"] = inspect command_line_config["use_venv"] = use_venv + if modify_config_file: + command_line_config["modify_config_file"] = modify_config_file command_line_config["original_command"] = original_command.strip() command_line_config["started_from"] = os.getcwd() diff --git a/esm_runscripts/compute.py b/esm_runscripts/compute.py index 775c04e..afae02b 100644 --- a/esm_runscripts/compute.py +++ b/esm_runscripts/compute.py @@ -6,6 +6,7 @@ import six import yaml from esm_calendar import Date +from colorama import Fore, Back, Style, init import esm_tools @@ -153,9 +154,9 @@ def modify_namelists(config): if config["general"]["verbose"]: six.print_("\n" "- Setting up namelists for this run...") - for model in config["general"]["valid_model_names"]: - six.print_("-" * 80) - six.print_("* %s" % config[model]["model"], "\n") + for index, model in enumerate(config["general"]["valid_model_names"]): + print(f'{index+1}) {config[model]["model"]}') + print() for model in config["general"]["valid_model_names"]: config[model] = Namelist.nmls_load(config[model]) @@ -168,13 +169,12 @@ def modify_namelists(config): ) if config["general"]["verbose"]: - print("end of namelist section") + print("::: end of namelist section\n") return config def copy_files_to_thisrun(config): if config["general"]["verbose"]: - six.print_("=" * 80, "\n") six.print_("PREPARING EXPERIMENT") # Copy files: six.print_("\n" "- File lists populated, proceeding with copy...") @@ -191,7 +191,6 @@ def copy_files_to_thisrun(config): def copy_files_to_work(config): if config["general"]["verbose"]: - six.print_("=" * 80, "\n") six.print_("PREPARING WORK FOLDER") config = copy_files( config, config["general"]["in_filetypes"], source="thisrun", target="work" @@ -330,6 +329,16 @@ def strip_python_tags(s): config_file.write(out) return config +def color_diff(diff): + for line in diff: + if line.startswith('+'): + yield Fore.GREEN + line + Fore.RESET + elif line.startswith('-'): + yield Fore.RED + line + Fore.RESET + elif line.startswith('^'): + yield Fore.BLUE + line + Fore.RESET + else: + yield line def update_runscript(fromdir, scriptsdir, tfile, gconfig, file_type): """ @@ -382,7 +391,7 @@ def update_runscript(fromdir, scriptsdir, tfile, gconfig, file_type): f"{fromdir + '/' + tfile} differs from " + f"{scriptsdir + '/' + tfile}:\n" ) - for line in difflib.unified_diff(script_t, script_o): + for line in color_diff(difflib.unified_diff(script_t, script_o)): differences += line # If the --update flag is used, notify that the target script will @@ -474,21 +483,14 @@ def copy_tools_to_thisrun(config): # In case there is no esm_tools or namelists in the experiment folder, # copy from the default esm_tools path if not os.path.isdir(tools_dir): - if config['general'].get("use_venv") or esm_rcfile.FUNCTION_PATH.startswith("NONE_YET"): - if config["general"]["verbose"]: - print("Copying standard yamls from: package interal configs") - esm_tools.copy_config_folder(tools_dir) - else: - if config["general"]["verbose"]: - print("Copying from: ", esm_rcfile.FUNCTION_PATH) - shutil.copytree(esm_rcfile.FUNCTION_PATH, tools_dir) + print("Copying standard yamls from: ", esm_rcfile.EsmToolsDir("FUNCTION_PATH")) + esm_tools.copy_config_folder(tools_dir) if not os.path.isdir(namelists_dir): - if esm_rcfile.get_rc_entry("NAMELIST_PATH", "NONE_YET").startswith("NONE_YET"): - if config["general"]["verbose"]: - print("Copying standard namelists from: package internal namelists") - esm_tools.copy_namelist_folder(namelists_dir) - else: - shutil.copytree(esm_rcfile.get_rc_entry("NAMELIST_PATH"), namelists_dir) + print( + "Copying standard namelists from: ", + esm_rcfile.EsmToolsDir("NAMELIST_PATH"), + ) + esm_tools.copy_namelist_folder(namelists_dir) # If ``fromdir`` and ``scriptsdir`` are the same, this is already a computing # simulation which means we want to use the script in the experiment folder, diff --git a/esm_runscripts/filelists.py b/esm_runscripts/filelists.py index 646595f..3b9c8d4 100644 --- a/esm_runscripts/filelists.py +++ b/esm_runscripts/filelists.py @@ -147,14 +147,33 @@ def complete_sources(config): def reuse_sources(config): if config["general"]["run_number"] == 1: return config - for filetype in config["general"]["reusable_filetypes"]: + + # MA: the changes below are to be able to specify model specific reusable_filetypes + # without changing the looping order (a model loop nested inside a file-type loop) + + # Put together all the possible reusable file types + all_reusable_filetypes = [] + for model in config["general"]["valid_model_names"] + ["general"]: + all_reusable_filetypes = list( + set(all_reusable_filetypes + config[model].get("reusable_filetypes", [])) + ) + # Loop through all the reusable file types + for filetype in all_reusable_filetypes: for model in config["general"]["valid_model_names"] + ["general"]: - if filetype + "_sources" in config[model]: + # Get the model-specific reusable_filetypes, if not existing, get the + # general ones + model_reusable_filetypes = config[model].get( + "reusable_filetypes", + config["general"]["reusable_filetypes"] + ) + # If _sources dictionary exists and filetype is in the + # model-specific filetype list then add the sources + if filetype + "_sources" in config[model] and filetype in model_reusable_filetypes: for categ in config[model][filetype + "_sources"]: config[model][filetype + "_sources"][categ] = ( config[model]["experiment_" + filetype + "_dir"] + "/" - + config[model][filetype + "_targets"][categ].split("/")[-1] + + config[model][filetype + "_targets"][categ] ) return config @@ -336,6 +355,20 @@ def assemble_intermediate_files_and_finalize_targets(config): return config +def find_valid_year(config, year): + for entry in config: + min_val = -50000000000 + max_val = 500000000000 + + from_info = float(config[entry].get("from", min_val)) + to_info = float(config[entry].get("to", max_val)) + + if from_info <= year <= to_info: + return entry + # if the current model year is out of the valid bounds, report and exit + print(f"Sorry, no entry found for year {year} in config {config}") + sys.exit(-1) + def replace_year_placeholder(config): for filetype in config["general"]["all_model_filetypes"]: for model in config["general"]["valid_model_names"] + ["general"]: @@ -345,77 +378,220 @@ def replace_year_placeholder(config): filetype + "_additional_information" ]: if file_category in config[model][filetype + "_targets"]: + + all_years = [config["general"]["current_date"].year] + + if ( + "need_timestep_before" + in config[model][ + filetype + "_additional_information" + ][file_category] + ): + all_years.append( + config["general"]["prev_date"].year + ) if ( - "@YEAR@" + "need_timestep_after" + in config[model][ + filetype + "_additional_information" + ][file_category] + ): + all_years.append( + config["general"]["next_date"].year + ) + if ( + "need_year_before" + in config[model][ + filetype + "_additional_information" + ][file_category] + ): + all_years.append( + config["general"]["current_date"].year - 1 + ) + if ( + "need_year_after" + in config[model][ + filetype + "_additional_information" + ][file_category] + ): + all_years.append( + config["general"]["current_date"].year + 1 + ) + + + if "need_2years_before" in config[model][filetype + "_additional_information"][file_category]: + all_years.append(config["general"]["current_date"].year - 2) + + if "need_2years_after" in config[model][filetype + "_additional_information" ][file_category]: + all_years.append(config["general"]["current_date"].year + 2) + + + all_years = list( + dict.fromkeys(all_years) + ) # removes duplicates + + # loop over all years (including year_before & after) + # change replace the @YEAR@ variable with the + # corresponding year + for year in all_years: + new_category = file_category + "_year_" + str(year) + + # if the source contains 'from' or 'to' information + # then they have a dict type + if type(config[model][filetype + "_sources"][file_category]) == dict: + + # process the 'from' and 'to' information in + # file sources and targets + config[model][filetype + "_sources"][new_category] = \ + find_valid_year( + config[model][filetype + "_sources"][file_category], + year + ) + + config[model][filetype + "_targets"][new_category] = \ + config[model][filetype + "_targets"][file_category] + + # replace @YEAR@ in the targets + if ( + "@YEAR@" + in config[model][filetype + "_targets"][new_category] + ): + new_target_name = config[model][ + filetype + "_targets" + ][new_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_targets"][ + new_category + ] = new_target_name + + # replace @YEAR@ in the sources + if ( + "@YEAR@" + in config[model][filetype + "_sources"][new_category] + ): + new_source_name = config[model][ + filetype + "_sources" + ][new_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_sources"][ + new_category + ] = new_source_name + + # value is not a dictionary. Ie. it does not + # have `from` or `to` attributes. This else + # block preserves these sections in the config. + else: + # create `new_category` from `file_category` + config[model][filetype + "_sources"][new_category] = \ + config[model][filetype + "_sources"][file_category] + + config[model][filetype + "_targets"][new_category] = \ + config[model][filetype + "_targets"][file_category] + + # replace @YEAR@ in the targets + if ( + "@YEAR@" + in config[model][filetype + "_targets"][new_category] + ): + new_target_name = config[model][ + filetype + "_targets" + ][new_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_targets"][ + new_category + ] = new_target_name + + # replace @YEAR@ in the sources + if ( + "@YEAR@" + in config[model][filetype + "_sources"][new_category] + ): + new_source_name = config[model][ + filetype + "_sources" + ][new_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_sources"][ + new_category + ] = new_source_name + + # end if + # end of the for year loop + + # deniz: new additions for @YEAR_1850@ + # these are the Kinne aerosol files for the background + # aerosol concentration. They are needed for years + # 1849, 1850, and 1851. All these 3 files are the same + # and ECHAM needs them + if ("@YEAR_1850@" in config[model][filetype + "_targets"][file_category] ): - all_years = [config["general"]["current_date"].year] - - if ( - "need_timestep_before" - in config[model][ - filetype + "_additional_information" - ][file_category] - ): - all_years.append( - config["general"]["prev_date"].year - ) - if ( - "need_timestep_after" - in config[model][ - filetype + "_additional_information" - ][file_category] - ): - all_years.append( - config["general"]["next_date"].year - ) - if ( - "need_year_before" - in config[model][ - filetype + "_additional_information" - ][file_category] - ): - all_years.append( - config["general"]["current_date"].year - 1 - ) - if ( - "need_year_after" - in config[model][ - filetype + "_additional_information" - ][file_category] - ): - all_years.append( - config["general"]["current_date"].year + 1 - ) - - all_years = list( - dict.fromkeys(all_years) - ) # removes duplicates - - for year in all_years: - + # only target name is changed since source file is for a fixed year (1850) + for year in [1849, 1850, 1851]: new_category = file_category + "_year_" + str(year) + + # add the sources and targets to the config + config[model][filetype + "_sources"][new_category] = \ + config[model][filetype + "_sources"][file_category] + + config[model][filetype + "_targets"][new_category] = \ + config[model][filetype + "_targets"][file_category] + + # construct the file target and add this to the config new_target_name = config[model][ filetype + "_targets" - ][file_category].replace("@YEAR@", str(year)) - new_source_name = config[model][ - filetype + "_sources" - ][file_category].replace("@YEAR@", str(year)) - + ][new_category].replace("@YEAR_1850@", str(year)) + config[model][filetype + "_targets"][ new_category ] = new_target_name - config[model][filetype + "_sources"][ - new_category - ] = new_source_name - del config[model][filetype + "_targets"][file_category] - del config[model][filetype + "_sources"][file_category] + del config[model][filetype + "_sources"][file_category] + del config[model][filetype + "_targets"][file_category] + # end of if additonal information + + year = config["general"]["current_date"].year + + for file_category in config[model][filetype + "_targets"]: + + if type(config[model][filetype + "_sources"][file_category]) == dict: + config[model][filetype + "_sources"][file_category] = \ + find_valid_year( + config[model][filetype + "_sources"][file_category], + year + ) + if ( + "@YEAR@" + in config[model][filetype + "_targets"][file_category] + ): + new_target_name = config[model][ + filetype + "_targets" + ][file_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_targets"][ + file_category + ] = new_target_name + if ( + "@YEAR@" + in config[model][filetype + "_sources"][file_category] + ): + new_source_name = config[model][ + filetype + "_sources" + ][file_category].replace("@YEAR@", str(year)) + + config[model][filetype + "_sources"][ + file_category + ] = new_source_name + + # end of if filetype in target + # end of model loop + # end of filetype loop return config def log_used_files(config): + if config["general"]["verbose"]: + print("\n::: Logging used files") filetypes = config["general"]["relevant_filetypes"] for model in config["general"]["valid_model_names"] + ["general"]: with open( @@ -454,12 +630,11 @@ def log_used_files(config): + config[model][filetype + "_targets"][category] ) if config["general"]["verbose"]: - print( - "- " - + config[model][filetype + "_targets"][category] - + " : " - + config[model][filetype + "_sources"][category] - ) + print() + print((f'- source: ' + f'{config[model][filetype + "_sources"][category]}')) + print((f'- target: ' + f'{config[model][filetype + "_targets"][category]}')) flist.write("\n") flist.write("\n") flist.write(80 * "-") @@ -541,6 +716,8 @@ def resolve_symlinks(file_source): def copy_files(config, filetypes, source, target): + if config["general"]["verbose"]: + print("\n::: Copying files") successful_files = [] missing_files = {} @@ -566,26 +743,11 @@ def copy_files(config, filetypes, source, target): targetblock = config[model][filetype + "_" + text_target] for categ in sourceblock: file_source = os.path.normpath(sourceblock[categ]) - # NOTE(PG): This is a really, really, REALLY bad hack and it - # makes me physically ill to look at: - # NOTE(MA): The previous implementation was not able to include - # namelists that have no ``namelist`` in their name. This is a more - # general implementation but it enforces the use of the - # ``namelists`` list to be defined for each model with namelists. - namelist_candidates = ( - [item for item in config[model].get("namelists", [])] - + ["namelist"] - ) - isnamelist = any(map(file_source.__contains__, namelist_candidates)) - if source == "init": - if isnamelist and file_source.startswith("NONE_YET"): - file_source = esm_tools.get_namelist_filepath( - file_source.replace("NONE_YET/", "") - ) file_target = os.path.normpath(targetblock[categ]) if config["general"]["verbose"]: - print(f"source: {file_source}") - print(f" --> target: {file_target}") + print() + print(f"- source: {file_source}") + print(f"- target: {file_target}") if file_source == file_target: if config["general"]["verbose"]: print( @@ -602,7 +764,7 @@ def copy_files(config, filetypes, source, target): # (same as with ``mkdir -p >``) os.makedirs(dest_dir) if not os.path.isfile(file_source): - print(f"File not found: {file_source}...") + print(f"WARNING: File not found: {file_source}") missing_files.update({file_target: file_source}) continue if os.path.isfile(file_target) and filecmp.cmp( @@ -626,27 +788,28 @@ def copy_files(config, filetypes, source, target): if not "files_missing_when_preparing_run" in config["general"]: config["general"]["files_missing_when_preparing_run"] = {} if config["general"]["verbose"]: - six.print_("--- WARNING: These files were missing:") + six.print_("\n\nWARNING: These files were missing:") for missing_file in missing_files: - print(" - " + missing_file + ": " + missing_files[missing_file]) + print(f'- missing source: {missing_files[missing_file]}') + print(f'- missing target: {missing_file}') + print() config["general"]["files_missing_when_preparing_run"].update(missing_files) return config def report_missing_files(config): + # this list is populated by the ``copy_files`` function in filelists.py if "files_missing_when_preparing_run" in config["general"]: config = _check_fesom_missing_files(config) if not config["general"]["files_missing_when_preparing_run"] == {}: - six.print_(80 * "=") print("MISSING FILES:") for missing_file in config["general"]["files_missing_when_preparing_run"]: - print("-- " + missing_file + ": ") - print( - " --> " - + config["general"]["files_missing_when_preparing_run"][missing_file] - ) + print() + print(f'- missing source: {config["general"]["files_missing_when_preparing_run"][missing_file]}') + print(f'- missing target: {missing_file}') if not config["general"]["files_missing_when_preparing_run"] == {}: six.print_(80 * "=") + print() return config @@ -713,16 +876,18 @@ def get_method(movement): def complete_all_file_movements(config): + + mconfig = config["general"] + if "defaults.yaml" in mconfig: + if "per_model_defaults" in mconfig["defaults.yaml"]: + if "file_movements" in mconfig["defaults.yaml"]["per_model_defaults"]: + mconfig["file_movements"] = copy.deepcopy(mconfig["defaults.yaml"]["per_model_defaults"]["file_movements"]) + del mconfig["defaults.yaml"]["per_model_defaults"]["file_movements"] + config = create_missing_file_movement_entries(config) for model in config["general"]["valid_model_names"] + ["general"]: mconfig = config[model] - if model == "general": - if "defaults.yaml" in mconfig: - if "per_model_defaults" in mconfig["defaults.yaml"]: - if "file_movements" in mconfig["defaults.yaml"]["per_model_defaults"]: - mconfig["file_movements"] = mconfig["defaults.yaml"]["per_model_defaults"]["file_movements"] - del mconfig["defaults.yaml"]["per_model_defaults"]["file_movements"] if "file_movements" in mconfig: for filetype in config["general"]["all_model_filetypes"] + ["scripts", "unknown"]: if filetype in mconfig["file_movements"]: @@ -749,7 +914,13 @@ def complete_all_file_movements(config): def get_movement(config, model, filetype, source, target): if source == "init": - if config["general"]["run_number"] == 1 or filetype not in config["general"]["reusable_filetypes"]: + # Get the model-specific reusable_filetypes, if not existing, get the + # general ones + model_reusable_filetypes = config[model].get( + "reusable_filetypes", + config["general"]["reusable_filetypes"] + ) + if config["general"]["run_number"] == 1 or filetype not in model_reusable_filetypes: return config[model]["file_movements"][filetype]["init_to_exp"] else: return config[model]["file_movements"][filetype]["exp_to_run"] diff --git a/esm_runscripts/helpers.py b/esm_runscripts/helpers.py index 71c0d7b..9d8354e 100644 --- a/esm_runscripts/helpers.py +++ b/esm_runscripts/helpers.py @@ -28,32 +28,12 @@ def evaluate(config, job_type, recipe_name): % setup_name ) sys.exit(1) - if config["general"].get("use_venv"): - recipe = esm_tools.get_config_filepath("esm_software/esm_runscripts/esm_runscripts.yaml") - need_to_parse_recipe = True - plugins_bare = esm_tools.get_config_filepath( - "esm_software/esm_runscripts/esm_plugins.yaml" - ) - need_to_parse_plugins = True - elif esm_rcfile.FUNCTION_PATH.startswith("NONE_YET"): - recipe = esm_tools.get_config_filepath( - "esm_software/esm_runscripts/esm_runscripts.yaml" - ) - need_to_parse_recipe = True - plugins_bare = esm_tools.get_config_filepath( - "esm_software/esm_runscripts/esm_plugins.yaml" - ) - need_to_parse_plugins = True - else: - recipe = ( - esm_rcfile.FUNCTION_PATH - + "/esm_software/esm_runscripts/esm_runscripts.yaml" - ) - need_to_parse_recipe = True - plugins_bare = ( - esm_rcfile.FUNCTION_PATH + "/esm_software/esm_runscripts/esm_plugins.yaml" - ) - need_to_parse_plugins = True + + FUNCTION_PATH = esm_rcfile.EsmToolsDir("FUNCTION_PATH") + recipe = FUNCTION_PATH + "esm_software/esm_runscripts/esm_runscripts.yaml" + need_to_parse_recipe = True + plugins_bare = FUNCTION_PATH + "/esm_software/esm_runscripts/esm_plugins.yaml" + need_to_parse_plugins = True framework_recipe = esm_plugin_manager.read_recipe( recipe, {"job_type": job_type}, need_to_parse_recipe diff --git a/esm_runscripts/last_minute.py b/esm_runscripts/last_minute.py new file mode 100644 index 0000000..560d3db --- /dev/null +++ b/esm_runscripts/last_minute.py @@ -0,0 +1,76 @@ +import os +import copy +import esm_parser + +class last_minute_changes: + def __init__(self, config): + self.modify_config_file = config["general"].get("modify_config_file") + + if self.modify_config_file: + self.modify_config_file_abspath = os.path.abspath(self.modify_config_file) + self.modify_config = esm_parser.yaml_file_to_dict(self.modify_config_file_abspath) + + config["general"]["modify_config"] = copy.deepcopy(self.modify_config) + config["general"]["modify_config_file_abspath"] = self.modify_config_file_abspath + +# kh 27.11.20 "original command" is not available for esm_master (but for esm_runscripts) + if "original_command" in "general": + config["general"]["original_command"] = config["general"]["original_command"].replace( + self.modify_config_file, + self.modify_config_file_abspath + ) + + else: + self.modify_config_file_abspath = self.modify_config = None + + +def apply_last_minute_changes(config): + config["general"]["modify_config_memo"] = last_minute_changes(config) + + modify_config = config["general"]["modify_config_memo"].modify_config + + if modify_config: + settings = modify_config.get("build_and_run_modifications", {}).get("machine", {}).get("chooseable_settings") + _modify_config_with_settings(config, settings) + + settings = modify_config.get("build_only_modifications", {}).get("machine", {}).get("environment_settings") + _modify_config_with_settings(config, settings) + + settings = modify_config.get("run_only_modifications", {}).get("machine", {}).get("chooseable_settings") + _modify_config_with_settings(config, settings) + + settings = modify_config.get("run_only_modifications", {}).get("batch_system", {}).get("direct_settings") + _modify_config_with_settings(config, settings) + + return config + + +def restore_protected_last_minute_changes(config): + if config["general"]["modify_config_memo"]: + if config["general"]["modify_config_memo"].modify_config: + config["general"]["modify_config"] = config["general"]["modify_config_memo"].modify_config + del config["general"]["modify_config_memo"] + +# kh 26.11.20 + if "modify_config_memo" in config["general"]: # Entry could exist but be False + del config["general"]["modify_config_memo"] + + return config + + +def _modify_config_with_settings(config, settings): + if settings: + for k, v in settings.items(): + path_to_key = k.split(".") + entry = path_to_key.pop() + selected_config = config + for k2 in path_to_key: + selected_config = selected_config[k2] + if type(selected_config) == dict: + selected_config[entry] = v + elif type(selected_config) == list: + selected_config.append(entry + "=" + v) + + else: + raise ValueError("unexpected container type (neither dict nor list") + diff --git a/esm_runscripts/namelists.py b/esm_runscripts/namelists.py index cec788b..6d16342 100644 --- a/esm_runscripts/namelists.py +++ b/esm_runscripts/namelists.py @@ -207,8 +207,8 @@ def apply_echam_disturbance(config): else: disturbance_file = None if config["general"]["verbose"]: - print( - config["general"]["experiment_scripts_dir"] + print("WARNING: " + + config["general"]["experiment_scripts_dir"] + "/disturb_years.dat", "was not found", ) @@ -288,9 +288,12 @@ def nmls_output(mconfig): for nml_name, nml_obj in six.iteritems(mconfig.get("namelists", {})): all_nmls[nml_name] = nml_obj # PG: or a string representation? for nml_name, nml in all_nmls.items(): - six.print_("Final Contents of ", nml_name, ":") + message = f'\nFinal Contents of {nml_name}:' + six.print_(message) + six.print_(len(message) * '-') nml.write(sys.stdout) - six.print_("\n", 40 * "+ ") + print('-' * 80) + print(f'::: end of the contents of {nml_name}\n') return mconfig diff --git a/esm_runscripts/prepare.py b/esm_runscripts/prepare.py index 3061d90..8caff84 100644 --- a/esm_runscripts/prepare.py +++ b/esm_runscripts/prepare.py @@ -1,10 +1,32 @@ from . import helpers +import sys def run_job(config): helpers.evaluate(config, "prepare", "prepare_recipe") return config +def mini_resolve_variable_date_file(date_file, config): + while "${" in date_file: + pre, post = date_file.split("${", 1) + variable, post = post.split("}", 1) + if "." in variable: + variable_section, variable = variable.split(".") + answer = config[variable_section].get(variable) + else: + answer = config["general"].get(variable) + if not answer: + answer = config.get("env", {}).get(variable) + if not answer: + try: + assert (variable.startswith("env.") or variable.startswith("general.")) + except AssertionError: + print("The date file contains a variable which is not in the >>env<< or >>general<< section. This is not allowed!") + print(f"date_file = {date_file}") + sys.exit(1) + date_file = pre + answer + post + return date_file + def _read_date_file(config): import os @@ -18,6 +40,9 @@ def _read_date_file(config): + config["general"]["setup_name"] + ".date" ) + + date_file = mini_resolve_variable_date_file(date_file, config) + if os.path.isfile(date_file): logging.info("Date file read from %s", date_file) with open(date_file) as date_file: @@ -308,7 +333,7 @@ def set_overall_calendar(config): config["general"]["calendar"] = Calendar(0) return config - + def find_last_prepared_run(config): from esm_calendar import Date, Calendar import os @@ -332,8 +357,8 @@ def find_last_prepared_run(config): next_date = current_date.add(delta_date) end_date = next_date - (0, 0, 1, 0, 0, 0) - - datestamp = ( + + datestamp = ( current_date.format( form=9, givenph=False, givenpm=False, givenps=False ) @@ -442,7 +467,7 @@ def _add_all_folders(config): "config", "restart_in", ] - config["general"]["reusable_filetypes"] = ["bin", "src"] + config["general"]["reusable_filetypes"] = config["general"].get("reusable_filetypes", ["bin", "src"]) config["general"]["thisrun_dir"] = ( config["general"]["experiment_dir"] diff --git a/esm_runscripts/sim_objects.py b/esm_runscripts/sim_objects.py index fb0ba13..32e2eec 100644 --- a/esm_runscripts/sim_objects.py +++ b/esm_runscripts/sim_objects.py @@ -165,10 +165,8 @@ def distribute_per_model_defaults(self, config): def add_esm_runscripts_defaults_to_config(self, config): - if config['general'].get("use_venv") or esm_rcfile.FUNCTION_PATH.startswith("NONE_YET"): - path_to_file = esm_tools.get_config_filepath("esm_software/esm_runscripts/defaults.yaml") - else: - path_to_file = esm_rcfile.FUNCTION_PATH + "/esm_software/esm_runscripts/defaults.yaml" + FUNCTION_PATH = esm_rcfile.EsmToolsDir("FUNCTION_PATH") + path_to_file = FUNCTION_PATH + "/esm_software/esm_runscripts/defaults.yaml" default_config = esm_parser.yaml_file_to_dict(path_to_file) config["general"]["defaults.yaml"] = default_config config = self.distribute_per_model_defaults(config) diff --git a/esm_runscripts/slurm.py b/esm_runscripts/slurm.py index 90f43de..dcf8def 100644 --- a/esm_runscripts/slurm.py +++ b/esm_runscripts/slurm.py @@ -53,35 +53,86 @@ def get_jobid(): """ return os.environ.get("SLURM_JOB_ID") + def calc_requirements_multi_srun(self, config): + print("Paul was here...") + for run_type in list(config['general']['multi_srun']): + current_hostfile = self.path+"_"+run_type + print(f"Writing to: {current_hostfile}") + start_proc = 0 + end_proc = 0 + with open(current_hostfile, "w") as hostfile: + for model in config['general']['multi_srun'][run_type]['models']: + start_proc, start_core, end_proc, end_core = self.mini_calc_reqs( + config, model, hostfile, + start_proc, start_core, end_proc, end_core + ) + config['general']['multi_srun'][run_type]['hostfile'] = os.path.basename(current_hostfile) + + + @staticmethod + def mini_calc_reqs(config, model, hostfile, start_proc, start_core, end_proc, end_core): + if "nproc" in config[model]: + end_proc = start_proc + int(config[model]["nproc"]) - 1 + if "omp_num_threads" in config[model]: + end_core = start_core + int(config[model]["nproc"])*int(config[model]["omp_num_threads"]) - 1 + elif "nproca" in config[model] and "nprocb" in config[model]: + end_proc = start_proc + int(config[model]["nproca"])*int(config[model]["nprocb"]) - 1 + + # KH 30.04.20: nprocrad is replaced by more flexible + # partitioning using nprocar and nprocbr + if "nprocar" in config[model] and "nprocbr" in config[model]: + if config[model]["nprocar"] != "remove_from_namelist" and config[model]["nprocbr"] != "remove_from_namelist": + end_proc += config[model]["nprocar"] * config[model]["nprocbr"] + + else: + return start_proc, start_core, end_proc, end_core + + scriptfolder = config["general"]["thisrun_scripts_dir"] + "../work/" + if config["general"].get("taskset", False): + command = "./" + config[model]["execution_command_script"] + scriptname="script_"+model+".ksh" + with open(scriptfolder+scriptname, "w") as f: + f.write("#!/bin/ksh"+"\n") + f.write("export OMP_NUM_THREADS="+str(config[model]["omp_num_threads"])+"\n") + f.write(command+"\n") + os.chmod(scriptfolder+scriptname, 0o755) + + progname="prog_"+model+".sh" + with open(scriptfolder+progname, "w") as f: + f.write("#!/bin/sh"+"\n") + f.write("(( init = "+str(start_core)+" + $1 ))"+"\n") + f.write("(( index = init * "+str(config[model]["omp_num_threads"])+" ))"+"\n") + f.write("(( slot = index % "+str(config["computer"]["cores_per_node"])+" ))"+"\n") + f.write("echo "+model+" taskset -c $slot-$((slot + "+str(config[model]["omp_num_threads"])+" - 1"+"))"+"\n") + f.write("taskset -c $slot-$((slot + "+str(config[model]["omp_num_threads"])+" - 1)) ./script_"+model+".ksh"+"\n") + os.chmod(scriptfolder+progname, 0o755) + + if "execution_command" in config[model]: + command = "./" + config[model]["execution_command"] + elif "executable" in config[model]: + command = "./" + config[model]["executable"] + else: + return start_proc, start_core, end_proc, end_core + hostfile.write(str(start_proc) + "-" + str(end_proc) + " " + command + "\n") + start_proc = end_proc + 1 + start_core = end_core + 1 + return start_proc, start_core, end_proc, end_core + + def calc_requirements(self, config): """ Calculates requirements and writes them to ``self.path``. """ + if config['general'].get('multi_srun'): + self.calc_requirements_multi_srun(config) + return start_proc = 0 + start_core = 0 end_proc = 0 + end_core = 0 with open(self.path, "w") as hostfile: for model in config["general"]["valid_model_names"]: - if "nproc" in config[model]: - end_proc = start_proc + int(config[model]["nproc"]) - 1 - elif "nproca" in config[model] and "nprocb" in config[model]: - end_proc = start_proc + int(config[model]["nproca"])*int(config[model]["nprocb"]) - 1 - - # KH 30.04.20: nprocrad is replaced by more flexible - # partitioning using nprocar and nprocbr - if "nprocar" in config[model] and "nprocbr" in config[model]: - if config[model]["nprocar"] != "remove_from_namelist" and config[model]["nprocbr"] != "remove_from_namelist": - end_proc += config[model]["nprocar"] * config[model]["nprocbr"] - - else: - continue - if "execution_command" in config[model]: - command = "./" + config[model]["execution_command"] - elif "executable" in config[model]: - command = "./" + config[model]["executable"] - else: - continue - hostfile.write(str(start_proc) + "-" + str(end_proc) + " " + command + "\n") - start_proc = end_proc + 1 + start_proc, start_core, end_proc, end_core = self.mini_calc_reqs(config, model, hostfile, start_proc, start_core, end_proc, end_core) @staticmethod diff --git a/esm_runscripts/virtual_env_builder.py b/esm_runscripts/virtual_env_builder.py index 33735df..7ab3944 100644 --- a/esm_runscripts/virtual_env_builder.py +++ b/esm_runscripts/virtual_env_builder.py @@ -59,30 +59,100 @@ def _source_and_run_bin_in_venv(venv_context, command, shell): return subprocess.check_call(command, shell=shell) def _install_tools(venv_context, config): - #_run_bin_in_venv(venv_context, ['pip', 'install', 'git+https://github.com/esm-tools/esm_tools']) + """ + Installs the ESM-Tools packages for a virtual environment, taking into account + the user's specifications for editable packages and desired branches. + + To control which packages are installed in editable mode the user can add the + following to their runscript: + + .. code-block:: yaml + + general: + install__editable: True/False + + To control which package branch is installed (compatible with editable mode and + non-editable mode) the user can add the following to their runscript: + + .. code-block:: yaml + + general: + install__branch: + + Parameters + ---------- + venv_context : type + Some description + config : dict + Configuration dictionary for this run + """ + + # First installation of all packages, with the desired mode (editable/non-editable), + # and branches, together with all their dependencies + _install_tools_general(venv_context, config) + # Some packages, such as `esm_tools` have other packages as dependencies (i.e. + # `esm_parser`). In the previous step, if a package is installed for which a + # dependency is editable and/or branched, this dependency gets back to non-editable + # realese-branch. The following line resinstalls all the editable/branched packages + # again, this time without dependencies. + _install_tools_general(venv_context, config, deps=False) + +def _install_tools_general(venv_context, config, deps=True): + ''' + Actual installer of ESM-Tools packages for virtual environments. Used by + `_install_tools` method to correctly install packages with the user's requested + options for each package (editable/non-editable and branch). See `_install_tools` + documentation for more information. + + Parameters + ---------- + venv_context : type + Some description + config : dict + Configuration dictionary for this run + deps : bool + Boolean indicating whether dependencies should be installed or not + ''' + # Setup the --no-deps flag if necessary + if not deps: + no_deps_flag = ["--no-deps"] + else: + no_deps_flag = [] + # Loop through the esm_tools packages to be installed for tool in esm_tools_modules: + # Module info (url, editable install, branch...) url = f"https://github.com/esm-tools/{tool}" user_wants_editable = config["general"].get(f"install_{tool}_editable", False) user_wants_branch = config["general"].get(f"install_{tool}_branch") + # If the package is editable install it in /src/esm-tools/ if user_wants_editable: # Make sure the directory exists: src_dir = pathlib.Path(config['general']['experiment_dir'] + f"/src/esm-tools/{tool}") - src_dir.mkdir(parents=True, exist_ok=True) + if not src_dir.exists(): + src_dir.mkdir(parents=True, exist_ok=True) + # Select branch if necessary if user_wants_branch: branch_command = f" -b {user_wants_branch} " else: branch_command = "" - subprocess.check_call(f"git clone --quiet {branch_command} {url} {src_dir}", shell=True) - _run_bin_in_venv(venv_context, ["pip", "install", "-q", f"--find-links={os.environ.get('HOME')}/.cache/pip/wheels", "-e", src_dir]) - _run_bin_in_venv(venv_context, ["pip", "wheel", "-q", f"--wheel-dir={os.environ.get('HOME')}/.cache/pip/wheels", src_dir]) - else: + # Clone from git + if deps: + subprocess.check_call(f"git clone --quiet {branch_command} {url} {src_dir}", shell=True) + # Carry out the editable installation (with or without dependencies) + _run_bin_in_venv(venv_context, ["pip", "install", "-q", f"--find-links={os.environ.get('HOME')}/.cache/pip/wheels", "-e", src_dir] + no_deps_flag) + _run_bin_in_venv(venv_context, ["pip", "wheel", "-q", f"--wheel-dir={os.environ.get('HOME')}/.cache/pip/wheels", "-e", src_dir] + no_deps_flag) + # If the package is not editable then do a standard installation. + # Note: this step only runs with the `--no-deps` flag if the user has specified + # a branch, as this flags means also that is the second time passing through + # here, and we don't want to waste time on installing everything a second time + # if not necessary. + elif deps or (not deps and user_wants_branch): url = f"git+{url}" if user_wants_branch: url += f"@{user_wants_branch}" # NOTE(PG): We need the -U flag to ensure the branch is actually installed. - _run_bin_in_venv(venv_context, ["pip", "install", '-q', f"--find-links={os.environ.get('HOME')}/.cache/pip/wheels", "-U", url]) - _run_bin_in_venv(venv_context, ["pip", "wheel", '-q', f"--wheel-dir={os.environ.get('HOME')}/.cache/pip/wheels", url]) - + _run_bin_in_venv(venv_context, ["pip", "install", '-q', f"--find-links={os.environ.get('HOME')}/.cache/pip/wheels", "-U", url] + no_deps_flag) + _run_bin_in_venv(venv_context, ["pip", "wheel", '-q', f"--wheel-dir={os.environ.get('HOME')}/.cache/pip/wheels", url] + no_deps_flag) def _install_required_plugins(venv_context, config): required_plugins = [] @@ -109,15 +179,20 @@ def venv_bootstrap(config): config["general"]["command_line_config"]["use_venv"] = config["general"]["use_venv"] if config["general"].get("use_venv", False): if not in_virtualenv(): - print(f"Building virtual env, please be patient (this takes about 3 minutes)...") - start_time = datetime.datetime.now() venv_path = pathlib.Path(config['general']['experiment_dir']).joinpath('.venv_esmtools') - venv_context = _venv_create(venv_path) - _run_python_in_venv(venv_context, ['-m', 'pip', '-q', 'install', '-U', 'pip']) - _run_python_in_venv(venv_context, ['-m', 'pip', '-q', 'install', '-U', 'wheel']) - _install_tools(venv_context, config) - _install_required_plugins(venv_context, config) - print(f"...finished {datetime.datetime.now() - start_time}, restarting your job in the virtual env") + if venv_path.exists(): + print(f"{venv_path} already exists, reusing...") + venv_context = _EnvBuilder(with_pip=True).ensure_directories(venv_path) + else: + print(f"Building virtual env, please be patient (this takes about 3 minutes)...") + start_time = datetime.datetime.now() + venv_path = pathlib.Path(config['general']['experiment_dir']).joinpath('.venv_esmtools') + venv_context = _venv_create(venv_path) + _run_python_in_venv(venv_context, ['-m', 'pip', '-q', 'install', '-U', 'pip']) + _run_python_in_venv(venv_context, ['-m', 'pip', '-q', 'install', '-U', 'wheel']) + _install_tools(venv_context, config) + _install_required_plugins(venv_context, config) + print(f"...finished {datetime.datetime.now() - start_time}, restarting your job in the virtual env") sys.argv[0] = pathlib.Path(sys.argv[0]).name # NOTE(PG): This next line allows the job to restart itself in the # virtual environment. @@ -180,7 +255,10 @@ def _integorate_user_venv(config): choices=[ 'Run in virtualenv (You may set the flag `--contained-run` during your run call or set `general.use_venv: True`)', 'Run using default installation (You may set the flag `--open-run` during your run call or set `general.use_venv: False`)', + "Quit right now to adapt your runscript", ]).ask() # returns value of selection + if "Quit" in response: + sys.exit(0) config['general']['use_venv'] = "Run in virtualenv" in response user_confirmed = questionary.confirm("Are you sure?").ask() if "Run in virtualenv" in response: diff --git a/setup.cfg b/setup.cfg index e876150..44a2ce0 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 5.0.14 +current_version = 5.0.17 commit = True tag = True diff --git a/setup.py b/setup.py index 3dd8bdf..c09be7e 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ "esm_motd @ git+https://github.com/esm-tools/esm_motd.git", "psutil", "f90nml", + "colorama", "coloredlogs", "tqdm", "sqlalchemy", @@ -62,6 +63,6 @@ test_suite='tests', tests_require=test_requirements, url='https://github.com/dbarbi/esm_runscripts', - version="5.0.14", + version="5.0.17", zip_safe=False, )