MCPcopy
hub / github.com/dmlc/dgl / test_parmetis_wrapper

Function test_parmetis_wrapper

tests/tools/test_parmetis.py:187–246  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

185
186@unittest.skipIf(True, reason="mpi is not available in CI test framework.")
187def test_parmetis_wrapper():
188 with tempfile.TemporaryDirectory() as root_dir:
189 num_chunks = 2
190 graph_name = "mag240m"
191 g = create_chunked_dataset(root_dir, num_chunks)
192 all_ntypes = g.ntypes
193 all_etypes = g.etypes
194 num_constraints = len(all_ntypes) + 3
195 num_institutions = g.num_nodes("institution")
196 num_authors = g.num_nodes("author")
197 num_papers = g.num_nodes("paper")
198
199 # Trigger ParMETIS.
200 schema_file = os.path.join(root_dir, "chunked-data/metadata.json")
201 preproc_input_dir = os.path.join(root_dir, "chunked-data")
202 preproc_output_dir = os.path.join(
203 root_dir, "chunked-data/preproc_output_dir"
204 )
205 parmetis_output_file = os.path.join(
206 os.getcwd(), f"{graph_name}_part.{num_chunks}"
207 )
208 partitions_dir = os.path.join(root_dir, "chunked-data/partitions_dir")
209 hostfile = os.path.join(root_dir, "ip_config.txt")
210 with open(hostfile, "w") as f:
211 f.write("127.0.0.1\n")
212 f.write("127.0.0.1\n")
213
214 num_nodes = g.num_nodes()
215 num_edges = g.num_edges()
216 stats_file = f"{graph_name}_stats.txt"
217 with open(stats_file, "w") as f:
218 f.write(f"{num_nodes} {num_edges} {num_constraints}")
219
220 os.system(
221 f"python3 tools/distpartitioning/parmetis_wrapper.py "
222 f"--schema_file {schema_file} "
223 f"--preproc_input_dir {preproc_input_dir} "
224 f"--preproc_output_dir {preproc_output_dir} "
225 f"--hostfile {hostfile} "
226 f"--num_parts {num_chunks} "
227 f"--parmetis_output_file {parmetis_output_file} "
228 f"--partitions_dir {partitions_dir} "
229 )
230 print("Executing Done.")
231
232 ntype_count = {
233 "author": num_authors,
234 "paper": num_papers,
235 "institution": num_institutions,
236 }
237 for ntype_name in ["author", "paper", "institution"]:
238 fname = os.path.join(partitions_dir, f"{ntype_name}.txt")
239 print(fname)
240 assert os.path.isfile(fname)
241
242 # Load and check the partition ids in this file.
243 part_ids = np.loadtxt(fname)
244 assert part_ids.shape[0] == ntype_count[ntype_name]

Callers

nothing calls this directly

Calls 5

create_chunked_datasetFunction · 0.90
num_nodesMethod · 0.45
joinMethod · 0.45
writeMethod · 0.45
num_edgesMethod · 0.45

Tested by

no test coverage detected