Batch transition messages to workers/clients

[mrocklin]

Currently every call to Scheduler.transition potentially initiates a send of a message to workers and clients. This can be somewhat costly, even given the BatchedSend logic. There are a few ways that we could address this:

1. We could reduce the overhead of BatchedSend
2. We could batch things ourself in transitions

cc @jakirkham

[mrocklin]

For BatchedSend the relevant call is here

distributed/distributed/batched.py

Lines 130 to 142 in 2089e79

	def send(self, msg):
	"""Schedule a message for sending to the other side
	This completes quickly and synchronously
	"""
	if self.comm is not None and self.comm.closed():
	raise CommClosedError
	self.message_count += 1
	self.buffer.append(msg)
	# Avoid spurious wakeups if possible
	if self.next_deadline is None:
	self.waker.set()

I've made the closed call cheaper in #4453 , but there are still calls to increment a counter, append to a list, and possibly wake an Event (although this should be cheap in common case). If we can hide the check on a closed comm then increment + list.append seems cheap if we can cythonize. This would help out not only the transitions code, but BatchedSend usage more broadly.

[mrocklin]

Alternatively we could change how transitions works. We could accumulate a set of messages in subsequent calls to Scheduler.transition_foo, and then send them all at once.

Currently each Scheduler.transition_foo creates new dicts for worker_msgs and client_msgs, and then we send these. Instead, we might make two defaultdict(list) at a higher transitions call (one for each of workers and clients), and pass down these collections to the transition_foo methods to append to. At the end of the transitions call we would then bulk send many of these messages.

The transition call is also used on its own in the codebase. One way around this is to move the current transition method into a _transition method, and then implement the same comm handling logic in a new transition method, the same way we do in transitions, but now with the expectation of only a single operation.

[jakirkham]

Yeah actually was thinking we could create a TransitionMessage object, which would contain both dicts. Thinking we can do this with both Worker and Client messages. They could all derive from a common base class to make them easy to check for.

The benefit being we can annotate these so Cython converts them to C structs allowing fast access, easy type checking, the ability to take them as arguments to various functions (allowing those functions to be more easily Cythonized). Plus we can fast path serialization and deserialization for any of these message class.

Just a thought 🙂

[mrocklin]

I'm less excited about introducing a message class for this specific use case. This is an important performance improvement, but not important enough I think to add that kind of inertia. I'm open to switch to class-based messages, but if we do so then I think we need to do so holistically across the project, not just for this one case.

[jakirkham]

Exactly that would be the idea. As I see it, we have two problems: transitions are still slow (this is coming down to their inputs and outputs) and serialization is slow (the data going over the wire is a bit unstructured atm or at least that structure is not broadly visible and exploitable). These are actually two sides of the same problem and structured message objects would help solve both

[jakirkham]

Should add I don’t think these classes need to be complicated. The main goal is to identify what attributes are relevant to communicate and their types. I don’t think there is a need for more than say a constructor and a way to serialize/deserialize them (either with a method or function). So no need for other methods

[mrocklin]

Yeah, if we want to explore a more standardized/typed messageing system then I'm open to that. I think that that would have to be a larger effort though that looked at all of the messages that we send, categorized them, and then proposed a broad change. If you want to take this on I'm all for it. I know that others like @jcrist have expressed support for this idea in the past.

For the specific objective of accelerating the scheduler though this seems like a large, but high-hanging fruit. I think that there are probably a few easier things to do in the meantime, notably your idea of trying to reuse msgpack to remove the need for extract_serialize, and also maybe batching together messages in transitions.

[jakirkham]

Agree standardizing the messaging system may be a bit more work. Though I think we will ultimately want to move in that direction to make more of the Scheduler transitions amenable to Cythonization (currently things like talking **kwargs limit this). Perhaps we can take a baby step in this direction by standardizing the input and output arguments of transition_* functions themselves. WDYT?

[jakirkham]

Ofc transitions and transition both call transition, which is fine. Though there are a few functions that call transition directly:

• add_worker
• remove_worker
• stimulus_task_finished
• stimulus_task_erred

Is there some reason these need to call transition? Could they call transitions instead? Or should they all handle communication as well? Any other ideas we should consider? Thoughts? 🙂

[mrocklin]

Yeah, I was looking at this. They certainly shouldn't all handle communication (that sounds painful to maintain). In some cases they specify keywords that are specific to the transition_foo function. It feels awkward to push those keywords through transitions.

…

On Mon, Jan 25, 2021 at 1:21 PM jakirkham ***@***.***> wrote: Ofc transitions and transition both call transition, which is fine. Though there are a few functions that call transition directly: - add_worker - remove_worker - stimulus_task_finished - stimulus_task_erred Is there some reason these need to call transition? Could they call transitions instead? Or should they all handle communication as well? Thoughts? 🙂 — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#4454 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTHHCBND32RKQDSA743S3XOETANCNFSM4WP2UMHA> .

[jakirkham]

Yep that makes sense. It might be possible to get them to call transitions (though don't know if I'm overlooking something here). Happy to give that a try if that sounds reasonable (or do something else if not).

Yeah the keyword arguments are a little awkward atm. Pushing them into a very simple class (in Cython a struct) might be one way to handle them. Maybe there are other ways.

[mrocklin]

I would keep all of the calls to the self.transition( method. I would rewrite that function so that it calls a new _transition method, and then handles communication.

Maybe something like the following?

def transitions(...):
    worker_msgs = defaultdict(list)
    client_msgs = defaultdict(list)
    for ....
        recommendations = self._transition(..., worker_msgs, client_msgs)

    for worker, messages in worker_msgs.items():
        send(worker, messages)
    for client, messages in client_msgs.items():
        send(client, messages)

def transition(...):
    worker_msgs = defaultdict(list)
    client_msgs = defaultdict(list)
 
   recommendations = self._transition(..., worker_msgs, client_msgs)

    for worker, messages in worker_msgs.items():
        send(worker, messages)
    for client, messages in client_msgs.items():
        send(client, messages)

def _transition(...):
    # current definition of self.transition

[jakirkham]

I think PR ( #4451 ) is now doing a reasonable job of moving communication to transitions. IIUC this is basically at the end of handling an event on the Scheduler. Though please correct me if I'm misunderstanding that.

Did notice that add_worker calls transitions a couple of times. Not sure if there is an easy way to reduce that, but we may want to consider if that is an option. All other functions seem to call transitions once.

If we are able to avoid **kwargs in the transition and transition_* functions, we should be able to add C APIs for them, which should cutdown on the call time overhead.

[jakirkham]

Something else I'm curious about. What happens if send a list of messages to the Worker or Client? Do they know how to the process a list? Or do they only handle dicts? Wondering if we can combine all messages to an endpoint in one send.