code: def test_parmetis_wrapper():
      with tempfile.TemporaryDirectory() as root_dir:
          num_chunks = 2
          graph_name = "mag240m"
          g = create_chunked_dataset(root_dir, num_chunks)
          all_ntypes = g.ntypes
          all_etypes = g.etypes
          num_constraints = len(all_ntypes) + 3
          num_institutions = g.num_nodes("institution")
          num_authors = g.num_nodes("author")
          num_papers = g.num_nodes("paper")
  
          # Trigger ParMETIS.
          schema_file = os.path.join(
              root_dir, "chunked-data/metadata.json"
          )
          preproc_input_dir = os.path.join(root_dir, "chunked-data")
          preproc_output_dir = os.path.join(
              root_dir, "chunked-data/preproc_output_dir"
          )
          parmetis_output_file = os.path.join(
              os.getcwd(), f"{graph_name}_part.{num_chunks}"
          )
          partitions_dir = os.path.join(
              root_dir, "chunked-data/partitions_dir"
          )
          hostfile = os.path.join(root_dir, "ip_config.txt")
          with open(hostfile, "w") as f:
              f.write("127.0.0.1\n")
              f.write("127.0.0.1\n")
  
          num_nodes = g.num_nodes()
          num_edges = g.num_edges()
          stats_file = f"{graph_name}_stats.txt"
          with open(stats_file, "w") as f:
              f.write(f"{num_nodes} {num_edges} {num_constraints}")
  
          os.system(
              f"python3 tools/distpartitioning/parmetis_wrapper.py "
              f"--schema_file {schema_file} "
              f"--preproc_input_dir {preproc_input_dir} "
              f"--preproc_output_dir {preproc_output_dir} "
              f"--hostfile {hostfile} "
              f"--num_parts {num_chunks} "
              f"--parmetis_output_file {parmetis_output_file} "
              f"--partitions_dir {partitions_dir} "
          )
          print("Executing Done.")
  
          ntype_count = {
              "author": num_authors,
              "paper": num_papers,
              "institution": num_institutions,
          }
          for ntype_name in ["author", "paper", "institution"]:
              fname = os.path.join(partitions_dir, f"{ntype_name}.txt")
              print(fname)
              assert os.path.isfile(fname)
  
              # Load and check the partition ids in this file.
              part_ids = np.loadtxt(fname)
              assert part_ids.shape[0] == ntype_count[ntype_name]
              assert np.min(part_ids) == 0
              assert np.max(part_ids) == (num_chunks - 1)
  
file_path: dgl/tests/tools/test_parmetis.py
imports: import argparse
  import json
  import os
  import sys
  import tempfile
  import unittest
  import dgl
  import numpy
  import torch
  from dgl.data.utils import load_graphs, load_tensors
  from partition_algo.base import load_partition_meta
  from pytest_utils import create_chunked_dataset
label_desc: Use image node_id 3 for calling a global function with example usage: test_parmetis_wrapper() without return types