def get_partition_ids(self, agg_global_nids):
        """
        This function is used to get the partition-ids for a given set of global node ids

        global_nids <-> partition-ids mappings are deterministically  distributed across
        all the participating processes, within the service. A contiguous global-nids
        (ntype-ids, per-type-nids) are stored within each process and this is determined
        by the total no. of nodes of a given ntype-id and the rank of the process.

        Process, where the global_nid <-> partition-id mapping is stored can be easily computed
        as described above. Once this is determined we perform an alltoallv to send the request.
        On the receiving side, each process receives a set of global_nids and retrieves corresponding
        partition-ids using locally stored lookup tables. It builds responses to all the other
        processes and performs alltoallv.

        Once the response, partition-ids, is received, they are re-ordered corresponding to the
        incoming global-nids order and returns to the caller.

        Parameters:
        -----------
        self : instance of this class
            instance of this class, which is passed by the runtime implicitly

        agg_global_nids : numpy array
            an array of aggregated global node-ids for which partition-ids are
            to be retrieved by the distributed lookup service.

        Returns:
        --------
        list of integers :
            list of integers, which are the partition-ids of the global-node-ids (which is the
            function argument)
        """
        CHUNK_SIZE = 200 * 1000 * 1000
        # Determine the no. of times each process has to send alltoall messages.
        local_rows = agg_global_nids.shape[0]
        all_sizes = allgather_sizes(
            [local_rows], self.world_size, self.num_parts, return_sizes=True
        )
        max_count = np.amax(all_sizes)

        if max_count <= 0:
            logging.debug(
                f"[Rank: {self.rank}] No process has global_nids to process !!!"
            )
            return

        num_splits = np.ceil(max_count / CHUNK_SIZE).astype(np.uint16)
        LOCAL_CHUNK_SIZE = np.ceil(local_rows / num_splits).astype(np.int64)
        agg_partition_ids = []

        logging.debug(
            f"[Rank: {self.rank}] BatchSize: {CHUNK_SIZE}, \
                            max_count: {max_count}, \
                            splits: {num_splits}, \
                            rows: {agg_global_nids.shape}, \
                            local batch_size: {LOCAL_CHUNK_SIZE}"
        )

        for split in range(num_splits):
            # Compute the global_nids for this iteration
            global_nids = agg_global_nids[
                split * LOCAL_CHUNK_SIZE : (split + 1) * LOCAL_CHUNK_SIZE
            ]

            # Find the process where global_nid --> partition-id(owner) is stored.
            if len(global_nids) > 0:
                ntype_ids, type_nids = self.id_map(global_nids)
                ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy()
            else:
                ntype_ids = np.array([], dtype=np.int64)
                type_nids = np.array([], dtype=np.int64)

            assert len(ntype_ids) == len(global_nids)

            # For each node-type, the per-type-node-id <-> partition-id mappings are
            # stored as contiguous chunks by this lookup service.
            # The no. of these mappings stored by each process, in the lookup service, are
            # equally split among all the processes in the lookup service, deterministically.
            typeid_counts = self.ntype_count[ntype_ids]
            chunk_sizes = np.ceil(typeid_counts / self.world_size).astype(
                np.int64
            )
            service_owners = np.floor_divide(type_nids, chunk_sizes).astype(
                np.int64
            )

            # Now `service_owners` is a list of ranks (process-ids) which own the corresponding
            # global-nid <-> partition-id mapping.

            # Split the input global_nids into a list of lists where each list will be
            # sent to the respective rank/process
            # We also need to store the indices, in the indices_list, so that we can re-order
            # the final result (partition-ids) in the same order as the global-nids (function argument)
            send_list = []
            indices_list = []
            for idx in range(self.world_size):
                idxes = np.where(service_owners == idx)
                ll = global_nids[idxes[0]]
                send_list.append(torch.from_numpy(ll))
                indices_list.append(idxes[0])
            assert len(np.concatenate(indices_list)) == len(global_nids)
            assert np.all(
                np.sort(np.concatenate(indices_list))
                == np.arange(len(global_nids))
            )

            # Send the request to everyone else.
            # As a result of this operation, the current process also receives a list of lists
            # from all the other processes.
            # These lists are global-node-ids whose global-node-ids <-> partition-id mappings
            # are owned/stored by the current process
            owner_req_list = alltoallv_cpu(
                self.rank, self.world_size, send_list
            )

            # Create the response list here for each of the request list received in the previous
            # step. Populate the respective partition-ids in this response lists appropriately
            out_list = []
            for idx in range(self.world_size):
                if owner_req_list[idx] is None:
                    out_list.append(torch.empty((0,), dtype=torch.int64))
                    continue
                # Get the node_type_ids and per_type_nids for the incoming global_nids.
                ntype_ids, type_nids = self.id_map(owner_req_list[idx].numpy())
                ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy()

                # Lists to store partition-ids for the incoming global-nids.
                type_id_lookups = []
                local_order_idx = []

                # Now iterate over all the node_types and acculumulate all the partition-ids
                # since all the partition-ids are based on the node_type order... they
                # must be re-ordered as per the order of the input, which may be different.
                for tid in range(len(self.partid_list)):
                    cond = ntype_ids == tid
                    local_order_idx.append(np.where(cond)[0])
                    global_type_nids = type_nids[cond]
                    if len(global_type_nids) <= 0:
                        continue

                    local_type_nids = (
                        global_type_nids - self.type_nid_begin[tid]
                    )

                    assert np.all(local_type_nids >= 0)
                    assert np.all(
                        local_type_nids
                        <= (
                            self.type_nid_end[tid]
                            + 1
                            - self.type_nid_begin[tid]
                        )
                    )

                    cur_owners = self.partid_list[tid][local_type_nids]
                    type_id_lookups.append(cur_owners)

                # Reorder the partition-ids, so that it agrees with the input order --
                # which is the order in which the incoming message is received.
                if len(type_id_lookups) <= 0:
                    out_list.append(torch.empty((0,), dtype=torch.int64))
                else:
                    # Now reorder results for each request.
                    sort_order_idx = np.argsort(np.concatenate(local_order_idx))
                    lookups = np.concatenate(type_id_lookups)[sort_order_idx]
                    out_list.append(torch.from_numpy(lookups))

            # Send the partition-ids to their respective requesting processes.
            owner_resp_list = alltoallv_cpu(
                self.rank, self.world_size, out_list
            )

            # Owner_resp_list, is a list of lists of numpy arrays where each list
            # is a list of partition-ids which the current process requested
            # Now we need to re-order so that the parition-ids correspond to the
            # global_nids which are passed into this function.

            # Order according to the requesting order.
            # Owner_resp_list is the list of owner-ids for global_nids (function argument).
            owner_ids = [x for x in owner_resp_list if x is not None]
            if len(owner_ids) > 0:
                owner_ids = torch.cat(owner_ids).numpy()
            else:
                owner_ids = np.array([], dtype=np.int64)
            assert len(owner_ids) == len(global_nids)

            global_nids_order = np.concatenate(indices_list)
            sort_order_idx = np.argsort(global_nids_order)
            owner_ids = owner_ids[sort_order_idx]
            global_nids_order = global_nids_order[sort_order_idx]
            assert np.all(np.arange(len(global_nids)) == global_nids_order)

            if len(owner_ids) > 0:
                # Store the partition-ids for the current split
                agg_partition_ids.append(owner_ids)

        # Stitch the list of partition-ids and return to the caller
        if len(agg_partition_ids) > 0:
            agg_partition_ids = np.concatenate(agg_partition_ids)
        else:
            agg_partition_ids = np.array([], dtype=np.int64)
        assert agg_global_nids.shape[0] == agg_partition_ids.shape[0]

        # Now the owner_ids (partition-ids) which corresponding to the  global_nids.
        return agg_partition_ids

        This function is used to get the partition-ids for a given set of global node ids

        global_nids <-> partition-ids mappings are deterministically  distributed across
        all the participating processes, within the service. A contiguous global-nids
        (ntype-ids, per-type-nids) are stored within each process and this is determined
        by the total no. of nodes of a given ntype-id and the rank of the process.

        Process, where the global_nid <-> partition-id mapping is stored can be easily computed
        as described above. Once this is determined we perform an alltoallv to send the request.
        On the receiving side, each process receives a set of global_nids and retrieves corresponding
        partition-ids using locally stored lookup tables. It builds responses to all the other
        processes and performs alltoallv.

        Once the response, partition-ids, is received, they are re-ordered corresponding to the
        incoming global-nids order and returns to the caller.

        Parameters:
        -----------
        self : instance of this class
            instance of this class, which is passed by the runtime implicitly

        agg_global_nids : numpy array
            an array of aggregated global node-ids for which partition-ids are
            to be retrieved by the distributed lookup service.

        Returns:
        --------
        list of integers :
            list of integers, which are the partition-ids of the global-node-ids (which is the
            function argument)
        
tools/distpartitioning/dist_lookup.py