()
| 115 | |
| 116 | |
| 117 | def test_parmetis_postprocessing(): |
| 118 | with tempfile.TemporaryDirectory() as root_dir: |
| 119 | num_chunks = 2 |
| 120 | g = create_chunked_dataset(root_dir, num_chunks) |
| 121 | |
| 122 | num_nodes = g.num_nodes() |
| 123 | num_institutions = g.num_nodes("institution") |
| 124 | num_authors = g.num_nodes("author") |
| 125 | num_papers = g.num_nodes("paper") |
| 126 | |
| 127 | # Generate random parmetis partition ids for the nodes in the graph. |
| 128 | # Replace this code with actual ParMETIS executable when it is ready |
| 129 | output_dir = os.path.join(root_dir, "chunked-data") |
| 130 | assert os.path.isdir(output_dir) |
| 131 | |
| 132 | parmetis_file = os.path.join(output_dir, "parmetis_output.txt") |
| 133 | node_ids = np.arange(num_nodes) |
| 134 | partition_ids = np.random.randint(0, 2, (num_nodes,)) |
| 135 | parmetis_output = np.column_stack([node_ids, partition_ids]) |
| 136 | |
| 137 | # Create parmetis output, this is mimicking running actual parmetis. |
| 138 | with open(parmetis_file, "w") as f: |
| 139 | np.savetxt(f, parmetis_output) |
| 140 | assert os.path.isfile(parmetis_file) |
| 141 | |
| 142 | # Check the post processing script here. |
| 143 | results_dir = os.path.join(output_dir, "partitions_dir") |
| 144 | json_file = os.path.join(output_dir, "metadata.json") |
| 145 | print(json_file) |
| 146 | print(results_dir) |
| 147 | print(parmetis_file) |
| 148 | os.system( |
| 149 | f"python3 tools/distpartitioning/parmetis_postprocess.py " |
| 150 | f"--postproc_input_dir {output_dir} " |
| 151 | f"--schema_file metadata.json " |
| 152 | f"--parmetis_output_file {parmetis_file} " |
| 153 | f"--partitions_dir {results_dir}" |
| 154 | ) |
| 155 | |
| 156 | ntype_count = { |
| 157 | "author": num_authors, |
| 158 | "paper": num_papers, |
| 159 | "institution": num_institutions, |
| 160 | } |
| 161 | for ntype_name in ["author", "paper", "institution"]: |
| 162 | fname = os.path.join(results_dir, f"{ntype_name}.txt") |
| 163 | print(fname) |
| 164 | assert os.path.isfile(fname) |
| 165 | |
| 166 | # Load and check the partition ids in this file. |
| 167 | part_ids = np.loadtxt(fname) |
| 168 | assert part_ids.shape[0] == ntype_count[ntype_name] |
| 169 | assert np.min(part_ids) == 0 |
| 170 | assert np.max(part_ids) == 1 |
| 171 | |
| 172 | # check partition meta file |
| 173 | part_meta_file = os.path.join(results_dir, "partition_meta.json") |
| 174 | assert os.path.isfile(part_meta_file) |
nothing calls this directly
no test coverage detected