|
| 1 | +import sys |
| 2 | +import toml |
| 3 | +import os |
| 4 | + |
| 5 | +def read_cargo_toml(file_path): |
| 6 | + with open(file_path, "r") as cargo_toml_file: |
| 7 | + cargo_toml = toml.load(cargo_toml_file) |
| 8 | + return cargo_toml |
| 9 | + |
| 10 | +def get_package_name(crate_dir): |
| 11 | + cargo_toml_path = os.path.join(crate_dir, "Cargo.toml") |
| 12 | + cargo_toml = read_cargo_toml(cargo_toml_path) |
| 13 | + |
| 14 | + if cargo_toml is not None: |
| 15 | + package = cargo_toml.get("package", {}) |
| 16 | + return package.get("name", None) |
| 17 | + else: |
| 18 | + raise Exception("Cargo.toml not found for crate {}.".format(crate_dir)) |
| 19 | + |
| 20 | +def group_crates_by_first_directory(members): |
| 21 | + ''' |
| 22 | + Here we ensure that all crates in the same directory are grouped together, to minimize fracturing of coverage. |
| 23 | + We assume here that crates are tested within their directory. This is a fair assumption since we put test-suites |
| 24 | + in the same directory as the crate they are testing. |
| 25 | + ''' |
| 26 | + crate_groups = {} |
| 27 | + for crate_dir in members: |
| 28 | + # Split the path name with the directory separator |
| 29 | + dir_parts = crate_dir.split(os.path.sep) |
| 30 | + if dir_parts: |
| 31 | + # Group crates by the first element and join them back using the separator |
| 32 | + dir_name = os.path.sep.join(dir_parts[:1]) |
| 33 | + if dir_name not in crate_groups: |
| 34 | + crate_groups[dir_name] = [] |
| 35 | + crate_groups[dir_name].append(crate_dir) |
| 36 | + return crate_groups |
| 37 | + |
| 38 | +def partition_workspace(m, n): |
| 39 | + cargo_toml = read_cargo_toml("Cargo.toml") |
| 40 | + |
| 41 | + if cargo_toml is None: |
| 42 | + raise Exception("Cargo.toml not found.") |
| 43 | + |
| 44 | + workspace = cargo_toml.get("workspace", {}) |
| 45 | + members = workspace.get("members", []) |
| 46 | + |
| 47 | + if not members: |
| 48 | + raise Exception("No members found in the workspace.") |
| 49 | + |
| 50 | + # Group crates based on directory structure using directory separators |
| 51 | + crate_directory_groups = group_crates_by_first_directory(members) |
| 52 | + |
| 53 | + # Calculate elements per partition and remainder based on the number of crate directories |
| 54 | + total_directories = len(crate_directory_groups) |
| 55 | + elements_per_partition = total_directories // m |
| 56 | + remainder = total_directories % m |
| 57 | + |
| 58 | + # Calculate the start and end indices for the current partition |
| 59 | + start_idx = n * elements_per_partition + min(n, remainder) |
| 60 | + end_idx = start_idx + elements_per_partition + (1 if n < remainder else 0) |
| 61 | + |
| 62 | + # Get the crate directories for the current partition |
| 63 | + partition_directories = [] |
| 64 | + for dir_path in list(crate_directory_groups.keys())[start_idx:end_idx]: |
| 65 | + partition_directories.extend(crate_directory_groups[dir_path]) |
| 66 | + |
| 67 | + # Get the package names from the crate directories in this partition |
| 68 | + package_names = [get_package_name(crate_dir) for crate_dir in partition_directories] |
| 69 | + |
| 70 | + return package_names |
| 71 | + |
| 72 | +if __name__ == "__main__": |
| 73 | + if len(sys.argv) == 2 and (sys.argv[1] == "--help" or sys.argv[1] == "-h"): |
| 74 | + print("Usage: python partition_workspace.py <total_partitions> <partition_index> [prefix]") |
| 75 | + print("") |
| 76 | + print("Partitions the workspace into m partitions and returns the crates in the n-th partition.") |
| 77 | + print("This can be used to split the workload of running tests across multiple CI jobs.") |
| 78 | + print("") |
| 79 | + print("To run the tests for the n-th partition, use the following command (for 3 partitions):") |
| 80 | + print("---") |
| 81 | + print("python3 {} 3 0 -p | xargs cargo test".format(sys.argv[0])) |
| 82 | + print("python3 {} 3 1 -p | xargs cargo test".format(sys.argv[0])) |
| 83 | + print("python3 {} 3 2 -p | xargs cargo test".format(sys.argv[0])) |
| 84 | + print("---") |
| 85 | + sys.exit(1) |
| 86 | + |
| 87 | + if len(sys.argv) < 3 or len(sys.argv) > 4: |
| 88 | + print("Usage: python3 {} <total_partitions> <partition_index> [prefix]".format(sys.argv[0])) |
| 89 | + sys.exit(1) |
| 90 | + |
| 91 | + total_splits = int(sys.argv[1]) |
| 92 | + partition_index = int(sys.argv[2]) |
| 93 | + prefix = sys.argv[3] + " " if len(sys.argv) == 4 else "" |
| 94 | + |
| 95 | + if total_splits <= 0 or partition_index < 0 or partition_index >= total_splits: |
| 96 | + print("Invalid input.") |
| 97 | + sys.exit(1) |
| 98 | + |
| 99 | + partition = partition_workspace(total_splits, partition_index) |
| 100 | + partition_with_prefix = [prefix + member for member in partition] |
| 101 | + |
| 102 | + print("{}".format(" ".join(partition_with_prefix))) |
0 commit comments