async def METHOD_NAME(self, revoked: Set[TP], newly_assigned: Set[TP]) -> None:
    """Call when a rebalance is needed."""
    T = traced_from_parent_span()
    # for isolated_partitions agents we stop agents for revoked
    # partitions.
    for agent, tps in self._collect_agents_for_update(revoked).items():
        await T(agent.on_partitions_revoked)(tps)
    # for isolated_partitions agents we start agents for newly
    # assigned partitions
    for agent, tps in T(self._collect_agents_for_update)(newly_assigned).items():
        await T(agent.on_partitions_assigned)(tps)