idds/tests/group_submission.py

0001 import argparse
0002 import yaml
0003 import os
0004 import pathlib
0005 import subprocess
0006
0007
0008 class SafeDict(dict):
0009     def __missing__(self, key):
0010         return '{' + key + '}'
0011
0012
0013 def generate_new_files(src_file, dest_dir):
0014     # Load src_file
0015     with open(src_file, "r") as f:
0016         template = f.read()
0017
0018     with open(src_file, "r") as f:
0019         full_config = yaml.safe_load(f)
0020
0021     src_path = pathlib.Path(src_file)
0022
0023     group_params = full_config.get("groupParameters", [])
0024
0025     new_files = []
0026     if not group_params:
0027         new_file = os.path.join(dest_dir, f"{src_path.stem}_group_none.yaml")
0028         with open(new_file, "w") as f:
0029             f.write(template)
0030         new_files.append(new_file)
0031     else:
0032         # Generate files
0033         for i, group in enumerate(group_params, 1):
0034             filled = template.format_map(SafeDict(group))
0035             new_file = os.path.join(dest_dir, f"{src_path.stem}_group_{i}.yaml")
0036             with open(new_file, "w") as f:
0037                 f.write(filled)
0038             new_files.append(new_file)
0039     return new_files
0040
0041
0042 def generate_new_commands(new_files, command):
0043     new_commands = []
0044     if "{new_file}" in command:
0045         for new_file in new_files:
0046             temp_command = command
0047             new_command = temp_command.format_map(SafeDict({"new_file": new_file}))
0048             new_commands.append(new_command)
0049     else:
0050         for new_file in new_files:
0051             temp_command = command
0052             new_command = f"{temp_command} {new_file}"
0053             new_commands.append(new_command)
0054     return new_commands
0055
0056
0057 def main():
0058     parser = argparse.ArgumentParser()
0059     parser.add_argument("--src_file", type=str, required=True,
0060                         help="The source YAML file with template + groupParameters")
0061     parser.add_argument("--dest_dir", type=str, default="outputs",
0062                         help="Directory where generated files are stored")
0063     parser.add_argument("--command", type=str, required=True,
0064                         help="The command to run. If '{new_file}' is in the command, "
0065                              "it will be replaced. Otherwise, the new file path will be appended.")
0066     parser.add_argument("--run_command", action="store_true",
0067                         help="Only print the generated command without executing if not setting")
0068     args = parser.parse_args()
0069
0070     # Ensure output directory exists
0071     os.makedirs(args.dest_dir, exist_ok=True)
0072
0073     new_files = generate_new_files(args.src_file, args.dest_dir)
0074     new_commands = generate_new_commands(new_files, args.command)
0075
0076     if not args.run_command:
0077         print("[DRY RUN]:")
0078         for cmd in new_commands:
0079             print(f"{cmd}")
0080     else:
0081         for cmd in new_commands:
0082             print(f"[RUN] {cmd}")
0083             subprocess.run(cmd, shell=True, check=True)
0084
0085
0086 if __name__ == "__main__":
0087     """
0088     examples:
0089         python group_submission.py --src_file test_cloud_us_group.yaml --command "bps submit" --run_command
0090         python group_submission.py --src_file test_cloud_us_group.yaml --command "bps submit {new_file}" --run_command
0091     """
0092     main()