class_docstr: PipelineStageManager is a helper class to manage pipeline stages.
  
  Args:
      pg_mesh (ProcessGroupMesh): Process group mesh.
      pipeline_axis (int): The axis along which the pipeline is constructed.
      is_virtual (bool): Whether to use circle p2p communication, it will make the first and last stage communicate with each other.
  
  Attributes:
      num_stages (int): Number of stages in the pipeline.
      stage (int): The current stage.
class_docstr_tok: ["PipelineStageManager","is","a","helper","class","to","manage","pipeline","stages",".","Args",":","pg_mesh","(","ProcessGroupMesh",")",":","Process","group","mesh",".","pipeline_axis","(","int",")",":","The","axis","along","which","the","pipeline","is","constructed",".","is_virtual","(","bool",")",":","Whether","to","use","circle","p2p","communication",",","it","will","make","the","first","and","last","stage","communicate","with","each","other",".","Attributes",":","num_stages","(","int",")",":","Number","of","stages","in","the","pipeline",".","stage","(","int",")",":","The","current","stage","."]
docstr: Get the p2p process group between two ranks. The order of the two ranks does not matter.
  
  Args:
      first_rank (int): The first rank.
      second_rank (int): The second rank.
  
  Returns:
      ProcessGroup: P2P process group between the two ranks.
docstr_tok: ["Get","the","p2p","process","group","between","two","ranks",".","The","order","of","the","two","ranks","does","not","matter",".","Args",":","first_rank","(","int",")",":","The","first","rank",".","second_rank","(","int",")",":","The","second","rank",".","Returns",":","ProcessGroup",":","P2P","process","group","between","the","two","ranks","."]
code: def get_p2p_process_group(
      self, first_rank: int, second_rank: int
  ) -> ProcessGroup:
      """Get the p2p process group between two ranks. The order of the two ranks does not matter.
  
      Args:
          first_rank (int): The first rank.
          second_rank (int): The second rank.
  
      Returns:
          ProcessGroup: P2P process group between the two ranks.
      """
      if first_rank > second_rank:
          first_rank, second_rank = second_rank, first_rank
      return self.p2p_groups[(first_rank, second_rank)]
  
file_path: colossalai/colossalai/pipeline/stage_manager.py
imports: import contextlib
  from typing import Dict, List, Optional, Tuple
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.cluster import ProcessGroupMesh
label_desc: Use image node_id 9 for calling the PipelineStageManager obj's underlying member method code with example usage: obj.get_p2p_process_group(first_rank, second_rank) and returns: self