def __init__()

in summarize_from_feedback/model_layout.py [0:0]

• 24 lines of code
• 4 McCabe index (conditional complexity)

    def __init__(self, layout, my_rank):
        """Layout is a numpy array with replica, shard"""
        self.layout = layout
        self.my_rank = my_rank

        self.total_gpus = layout.size
        self.all_ranks = list(range(self.total_gpus))

        self.n_replicas, self.n_shards = layout.shape

        if self.n_shards == 4:
            print(
                "WARNING: Using n_shards == 4 is currently slow because we have not"
                "implemented an efficient ring following the [0,1,3,2] pattern"
            )

        ([replica_idx], [shard_idx]) = np.where(layout == my_rank)

        self.replica_idx = int(replica_idx)
        self.shard_idx = int(shard_idx)

        # Create convenient accessors
        self.dp_sibling_ranks = [replica[shard_idx] for replica in layout]
        self.mp_sibling_ranks = list(layout[replica_idx])

        self.ranks_in_my_replica = layout[replica_idx].flatten().tolist()

        self.is_in_first_replica = self.replica_idx == 0

        self.replica_root = self.ranks_in_my_replica[0]
        self.is_replica_root = self.replica_root == self.my_rank

        self.is_logging_rank = self.is_replica_root and self.replica_idx == 0

        self.ranks_on_my_node = list(
            range(math.floor(self.my_rank / 8) * 8, 8 + math.floor(self.my_rank / 8) * 8)
        )