code: def test_parmetis_postprocessing():
      with tempfile.TemporaryDirectory() as root_dir:
          num_chunks = 2
          g = create_chunked_dataset(root_dir, num_chunks)
  
          num_nodes = g.num_nodes()
          num_institutions = g.num_nodes("institution")
          num_authors = g.num_nodes("author")
          num_papers = g.num_nodes("paper")
  
          # Generate random parmetis partition ids for the nodes in the graph.
          # Replace this code with actual ParMETIS executable when it is ready
          output_dir = os.path.join(root_dir, "chunked-data")
          assert os.path.isdir(output_dir)
  
          parmetis_file = os.path.join(
              output_dir, "parmetis_output.txt"
          )
          node_ids = np.arange(num_nodes)
          partition_ids = np.random.randint(0, 2, (num_nodes,))
          parmetis_output = np.column_stack([node_ids, partition_ids])
  
          # Create parmetis output, this is mimicking running actual parmetis.
          with open(parmetis_file, "w") as f:
              np.savetxt(f, parmetis_output)
          assert os.path.isfile(parmetis_file)
  
          # Check the post processing script here.
          results_dir = os.path.join(output_dir, "partitions_dir")
          json_file = os.path.join(output_dir, "metadata.json")
          print(json_file)
          print(results_dir)
          print(parmetis_file)
          os.system(
              f"python3 tools/distpartitioning/parmetis_postprocess.py "
              f"--postproc_input_dir {output_dir} "
              f"--schema_file metadata.json "
              f"--parmetis_output_file {parmetis_file} "
              f"--partitions_dir {results_dir}"
          )
  
          ntype_count = {
              "author": num_authors,
              "paper": num_papers,
              "institution": num_institutions,
          }
          for ntype_name in ["author", "paper", "institution"]:
              fname = os.path.join(results_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) == 1
  
          # check partition meta file
          part_meta_file = os.path.join(
              results_dir, "partition_meta.json"
          )
          assert os.path.isfile(part_meta_file)
          part_meta = load_partition_meta(part_meta_file)
          assert part_meta.num_parts == 2
          assert part_meta.algo_name == "metis"
  
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 2 for calling a global function with example usage: test_parmetis_postprocessing() without return types