preferred leader election...

So, I'm running into the case where after issuing a rolling restart, with
controlled shutdown enabled, the last server restarted ends up without any
partitions that it's the leader of. This is more pronounced of course if I
have only 2 servers in the cluster (during testing). I presume it's kind
of a problem to have imbalance in the number of partitions each server is
leader of? (Or maybe not such a big deal, since the servers still need to
forward messages between themselves anyway?).

So, I was looking at the 'preferred leader election' tool here:

https://cwiki.apache.org/confluence/display/KAFKA/Replication+tools

Is there a way to programmatically invoke this. I was thinking it make
sense to do periodically (e.g. once an hour), and also after a restart.
But I'd like it to happen from within my app (which contains kafka),
rather than invoking an external tool.

Thanks,

Jason