Microsoft has just released the Azure in-VM metadata service that will let you know when a machine is going to be rebooted.

it would be great to tie this into the node maintenance functionality so that the node is automatically drained.

It might be useful to do this for AWS as well (similar to https://github.com/scopely/shudder), but that looks a lot harder.