code: def gen_parmetis_input_args(params, schema_map):
      """Function to create two input arguments which will be passed to the parmetis.
      first argument is a text file which has a list of node-weights files,
      namely parmetis-nfiles.txt, and second argument is a text file which has a
      list of edge files, namely parmetis_efiles.txt.
      ParMETIS uses these two files to read/load the graph and partition the graph
      With regards to the file format, parmetis_nfiles.txt uses the following format
      for each line in that file:
          <filename> <global_node_id_start> <global_node_id_end>(exclusive)
      While parmetis_efiles.txt just has <filename> in each line.
  
      Parameters:
      -----------
      params : argparser instance
          Instance of ArgParser class, which has all the input arguments passed to
          run this program.
      schema_map : json dictionary
          Dictionary object created after reading the graph metadata.json file.
      """
  
      # TODO: This makes the assumption that all node files have the same number of chunks
      ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(
          schema_map
      )
      type_nid_dict, ntype_gnid_offset = get_idranges(
          schema_map[constants.STR_NODE_TYPE],
          dict(
              zip(
                  schema_map[constants.STR_NODE_TYPE],
                  schema_map[constants.STR_NUM_NODES_PER_TYPE],
              )
          ),
      )
  
      # Check if <graph-name>_stats.txt exists, if not create one using metadata.
      # Here stats file will be created in the current directory.
      # No. of constraints, third column in the stats file is computed as follows:
      #   num_constraints = no. of node types + train_mask + test_mask + val_mask
      #   Here, (train/test/val) masks will be set to 1 if these masks exist for
      #   all the node types in the graph, otherwise these flags will be set to 0
      assert (
          constants.STR_GRAPH_NAME in schema_map
      ), "Graph name is not present in the json file"
      graph_name = schema_map[constants.STR_GRAPH_NAME]
      if not os.path.isfile(
          os.path.join(params.input_dir, f"{graph_name}_stats.txt")
      ):
          num_nodes = np.sum(
              schema_map[constants.STR_NUM_NODES_PER_TYPE]
          )
          num_edges = np.sum(
              schema_map[constants.STR_NUM_EDGES_PER_TYPE]
          )
          num_ntypes = len(schema_map[constants.STR_NODE_TYPE])
  
          num_constraints = num_ntypes
  
          with open(
              os.path.join(params.input_dir, f"{graph_name}_stats.txt"),
              "w",
          ) as sf:
              sf.write(f"{num_nodes} {num_edges} {num_constraints}")
  
      node_files = []
      outdir = Path(params.output_dir)
      os.makedirs(outdir, exist_ok=True)
      for ntype_id, ntype_name in ntid_ntype_map.items():
          global_nid_offset = ntype_gnid_offset[ntype_name][0, 0]
          total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][
              ntype_id
          ]
          per_rank_range = np.ones(
              (params.num_parts,), dtype=np.int64
          ) * (total_count // params.num_parts)
          for i in range(total_count % params.num_parts):
              per_rank_range[i] += 1
          tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
          tid_end = np.cumsum(per_rank_range)
          logging.info(f" tid-start = {tid_start}, tid-end = {tid_end}")
          logging.info(f" per_rank_range - {per_rank_range}")
  
          for part_idx in range(params.num_parts):
              local_tid_start = tid_start[part_idx]
              local_tid_end = tid_end[part_idx]
              out_file = os.path.join(
                  outdir,
                  "node_weights_{}_{}.txt".format(ntype_name, part_idx),
              )
              node_files.append(
                  (
                      out_file,
                      global_nid_offset + local_tid_start,
                      global_nid_offset + local_tid_end,
                  )
              )
  
      with open(
          os.path.join(params.output_dir, "parmetis_nfiles.txt"), "w"
      ) as parmetis_nf:
          for node_file in node_files:
              # format: filename global_node_id_start global_node_id_end(exclusive)
              parmetis_nf.write(
                  "{} {} {}\n".format(
                      node_file[0], node_file[1], node_file[2]
                  )
              )
  
      # Regenerate edge files here.
      # NOTE: The file names need to match the ones generated by gen_edge_files function
      edge_data = schema_map[constants.STR_EDGES]
      edge_files = []
      for etype_name, etype_info in edge_data.items():
          edge_data_files = etype_info[constants.STR_DATA]
          for edge_file_path in edge_data_files:
              out_file_name = Path(edge_file_path).stem.split(".")[0]
              out_file = os.path.join(
                  outdir,
                  etype_name,
                  "edges_{}.csv".format(out_file_name),
              )
              edge_files.append(out_file)
  
      with open(
          os.path.join(params.output_dir, "parmetis_efiles.txt"), "w"
      ) as parmetis_efile:
          for edge_file in edge_files:
              parmetis_efile.write("{}\n".format(edge_file))
  
file_path: dgl/tools/distpartitioning/parmetis_preprocess.py
imports: import argparse
  import logging
  import os
  import platform
  from pathlib import Path
  import array_readwriter
  import constants
  import numpy
  import pyarrow
  import pyarrow.csv
  from utils import generate_read_list, generate_roundrobin_read_list, get_idranges, get_node_types, read_json
label_desc: Use image node_id 5 for calling a global function with example usage: gen_parmetis_input_args(params, schema_map) without return types