code: def optimize_parts_lp(
      vector_size: int, bandwidths: np.ndarray, min_size: int = 0
  ) -> np.ndarray:
      """
      This method solves an optimization problem to minimize the total allreduce time.
      In butterfly all-reduce, each peer acts both as a "client" and as an "aggregator":
      * a "client" splits his local vector into shards and sends each shard to one peer, then downloads the average
      * an "aggregator" receives a certain part of vector components from all peers, aggregates and returns the average
  
      Peer i network load as a "client" = vector_size * (1 - fraction_assigned_to_peer_i)
      Peer i network load as an "aggregator" = vector_size * (group_size - 1) * fraction_assigned_to_peer_i
      Peer i total communication = vector_size * [1 + (group_size - 2) * fraction_assigned_to_peer_i]
      Total time = max_i (total_communication_for_peer_i / bandwidths[i])
  
      We solve this optimization problem by reducing it to linear programming with a minimax reduction
      (see lecture notes: https://www.usna.edu/Users/math/dphillip/sa305.s15/phillips/lessons/32/32.pdf )
  
      :returns: a vector of "scores", i-th score is proportional to the fraction of weights assigned to i-th peer
      """
      assert np.all(bandwidths >= 0) and np.any(bandwidths > 0)
      bandwidths = np.asarray(bandwidths, dtype=np.float64)
      permutation = np.argsort(-bandwidths)
      bandwidths = bandwidths[permutation]
      is_nonzero = bandwidths != 0
  
      group_size = len(bandwidths)
      num_variables = group_size + 1  # [w_1, ..., w_N, xi]
  
      c = np.zeros(num_variables, dtype=np.float64)
      c[-1] = 1.0  # optimize w.r.t. xi
  
      # the constraints below are tuples (A, b) such that Ax <= b
      nonnegative_weights = -np.eye(
          group_size, num_variables, dtype=c.dtype
      ), np.zeros(group_size, c.dtype)
      weights_sum_to_one = c[None, :] - 1.0, np.array([-1.0])
      coeff_per_variable = (group_size - 2.0) / np.maximum(
          bandwidths, 10**-LOAD_BALANCING_LP_DECIMALS
      )
      coeff_matrix_minus_xi = np.hstack(
          [
              np.diag(coeff_per_variable),
              -np.ones((group_size, 1), c.dtype),
          ]
      )
      xi_is_maximum = (
          coeff_matrix_minus_xi[is_nonzero],
          -1.0 / bandwidths[is_nonzero],
      )
      force_max_weights = np.eye(
          group_size, M=num_variables, dtype=c.dtype
      ), is_nonzero.astype(c.dtype)
  
      A, b = list(
          map(
              np.concatenate,
              zip(
                  nonnegative_weights,
                  weights_sum_to_one,
                  xi_is_maximum,
                  force_max_weights,
              ),
          )
      )
  
      solution = scipy.optimize.linprog(
          c, A_ub=A, b_ub=b, method="interior-point"
      )
      if solution.success:
          peer_scores = solution.x[:group_size]
          # if some peers have less than min_size elements, transfer their share to other peers (if any)
          if np.max(peer_scores) >= min_size / float(vector_size):
              peer_scores[
                  peer_scores < min_size / float(vector_size)
              ] = 0.0
          peer_scores = np.round(
              peer_scores, LOAD_BALANCING_LP_DECIMALS
          )
      else:
          logger.error(
              f"Failed to solve load-balancing for bandwidths {bandwidths}"
          )
          peer_scores = np.ones(group_size, c.dtype)
  
      return peer_scores[np.argsort(permutation)]
  
imports: from typing import Optional, Sequence, Tuple
  import numpy
  import scipy.optimize
  from hivemind.utils.logging import get_logger
label_desc: Use image node_id 2 for calling a global function with example usage: optimize_parts_lp(vector_size, bandwidths, min_size) and returns: peer_scores