Refactor SchedulerState from Scheduler

[jakirkham]

Requires PR ( #4343 )

This refactors out the SchedulerState as @cclass including info about endpoints and tasks. Also includes the transition_* functions. This should make it easier to more thoroughly Cythonize these components.

[jakirkham]

cc @mrocklin @quasiben

[jakirkham]

Should add that after doing this, I can see a path to typing all attributes of the Scheduler and using them throughout. This would speed up attribute access and would also give Cython the opportunity to use targeted Python C APIs with those attributes (assuming they have a builtin type more specific than object). This would involve refactoring these into a base class, say _Scheduler, and decorating it with @cclass, then going through and casting all usages of self to a new variable, say parent, typed as _Scheduler much as we have done here with SchedulerState.

[mrocklin]

FYI, I'm slammed today and tomorrow. I'm not going to be very responsive this week.

…

On Tue, Dec 15, 2020 at 7:35 AM jakirkham ***@***.***> wrote: Requires PR ( #4343 <#4343> ) This refactors out the SchedulerState as @cclass including info about endpoints and tasks. Also includes the transition_* functions. This should make it easier to more thoroughly Cythonize these components. ------------------------------ You can view, comment on, or merge this pull request online at: #4365 Commit Summary - Move `worker_send` into transition functions - Refactor `_task_to_msg` from `send_task_to_worker` - Move `report` out of `_add_to_memory` - Refactor out `_client_releases_keys` - Collect client recs in `_add_to_memory` - Use `_client_releases_keys` in transitions - Refactor out `_task_to_report_msg` - Collect and send worker messages from transitions - Handle `report` in `transition` - Add method to send a message to a specific client - Add `_task_to_client_msgs` - Replace `report_msg` with `client_msgs` - Create empty `SchedulerState` class - Move `transition*` methods into `SchedulerState` - Add attributes for `SchedulerState` - Initialize attributes in `SchedulerState` - Pass arguments to `super` class - Use `SchedulerState` attributes - Use `cast` to access parent class attributes - Drop no longer needed `cast`s & local assignments - Use `dict` views onto `SortedDict` where possible - Add ***@***.***`s for `SchedulerState` attributes File Changes - *M* distributed/core.py <https://github.com/dask/distributed/pull/4365/files#diff-581957c552b88dd04319efca1429c0c0827daff509a115542e70b661c4c04914> (3) - *M* distributed/scheduler.py <https://github.com/dask/distributed/pull/4365/files#diff-bbcf2e505bf2f9dd0dc25de4582115ee4ed4a6e80997affc7b22122912cc6591> (7323) Patch Links: - https://github.com/dask/distributed/pull/4365.patch - https://github.com/dask/distributed/pull/4365.diff — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#4365>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTA5GJL3PJPP5U43Q6DSU564HANCNFSM4U4QQSPA> .

[jakirkham]

No worries. Just wanted to make sure you are aware of this. Mostly just interested in your general high-level thoughts as opposed to a detailed review at this stage.

[jakirkham]

Some of the transition and supporting methods take **kwargs, but don't seem to do anything with them. As this hinders optimizing these methods, I've just dropped them. Please let me know if that is issue and they are needed by something else somewhere so we can devise alternative strategies (if needed).

[jakirkham]

Hmm...seems those **kwargs exist to just ignore some arguments (see traceback below)? Confused as to why they are showing up in the first place.

distributed.core - ERROR - transition_processing_memory() got an unexpected keyword argument 'status'
Traceback (most recent call last):
  File "/Users/jkirkham/Developer/distributed/distributed/core.py", line 592, in handle_stream
    handler(**merge(extra, msg))
  File "distributed/scheduler.py", line 4707, in distributed.scheduler.Scheduler.handle_task_finished
  File "distributed/scheduler.py", line 4115, in distributed.scheduler.Scheduler.stimulus_task_finished
  File "distributed/scheduler.py", line 6079, in distributed.scheduler.Scheduler.transition
  File "distributed/scheduler.py", line 6015, in distributed.scheduler.Scheduler.transition
  File "distributed/scheduler.py", line 2022, in distributed.scheduler.SchedulerState.transition_processing_memory
TypeError: transition_processing_memory() got an unexpected keyword argument 'status'

[jakirkham]

After some toying around, I think these are coming from here for finished tasks:

distributed/distributed/worker.py

Lines 1920 to 1929 in c2d8773

	d = {
	"op": "task-finished",
	"status": "OK",
	"key": ts.key,
	"nbytes": ts.nbytes,
	"thread": self.threads.get(ts.key),
	"type": typ_serialized,
	"typename": typename(typ),
	"metadata": ts.metadata,
	}

Also from errored tasks this appears to come from here:

distributed/distributed/worker.py

Lines 1931 to 1938 in c2d8773

	d = {
	"op": "task-erred",
	"status": "error",
	"key": ts.key,
	"thread": self.threads.get(ts.key),
	"exception": ts.exception,
	"traceback": ts.traceback,
	}

Perhaps the stimulus_* functions should be grabbing these extra arguments and handling them? Alternatively the transition_* functions getting these could list them explicitly (though they still are not used there). Or better still worker messages could skip passing info that is simply not needed. For now am just grabbing them in transition until we come to some agreement here on what to do.

[jakirkham]

Here's the Scheduler call graph with these changes so far. Some functions (like _add_to_memory) are being handled with a virtual table, which we don't need as they won't be inherited. So we might break them out into separate functions to make them independent of the class (and thus the virtual table).

Edit: We are spending less time in transition_waiting_processing (even relative to other transitions), which seems like a good thing. Also we are not spending as much time queueing stealable tasks. I'm wondering if part of the issue we saw with the Scheduler was due to it getting "antsy" and so trying things like letting tasks get stolen, which just created churn in turn slowing things down. Maybe we are starting to see a resolution to this behavior now?

[jakirkham]

Yeah messages including unused arguments and extra arguments being passed to transition_ functions seems like the main issue we need to address here (details in this comment #4365 (comment) and above it). Once we figure that out it should be more straightforward to optimize these.

[jakirkham]

There's other odd things like nbytes being set to None (-1 would be better for typing). worker being optional where a function absolutely needs it. And so on. Though these are second order problems we can fix if we fix the first order ones (above).

[madsbk]

Overall I like the refactoring; encapsulate the state of the scheduler in a different class makes a lot of sense. Actually, I think we should go even further and encapsulate all states of the scheduler and tasks with clearly defined invariants but that is future work :)

I have read through all of the changes and as far as I can see, it looks good but because of the amount of changes I am not confident that it is bug free :)

As a side note I agree, we should get rid of all the unused **kwargs in a future PR. It makes a harder to follow variables through the code.

[madsbk]

I am a bit confused, why do we need kwargs? And call to super?

[jakirkham]

Yeah that's fair. It's a little confusing.

The call to super is needed to call the parent constructors of Scheduler. These are ServerNode and SchedulerState. We also want to affect when these get called relative to other things in this constructor (as this constructor needs some of the attributes of the parent classes to be setup in later steps).

As Scheduler inherits from two base classes, we need to make sure that arguments for those constructors get passed through and picked up by the right parent class. There are some more details in this SO answer that may help clarify this further.

[madsbk]

I suggest we implement all the exception logging using a context manager as described in https://stackoverflow.com/a/28158006

[jakirkham]

Yeah I agree that something like that makes sense. Alternatively we might want to push this into transition where all of these methods are called. Held off on that for now as I didn't want to disrupt the code too much, but agree this is something we should follow up on.

[jakirkham]

Thanks for looking through this Mads! 😄

Overall I like the refactoring; encapsulate the state of the scheduler in a different class makes a lot of sense. Actually, I think we should go even further and encapsulate all states of the scheduler and tasks with clearly defined invariants but that is future work :)

Thanks Mads! 😄 Yeah for sure. I think this opens a lot of doors for us in terms of future optimizations. Just tried to find the right balance where we do enough to get some benefit without pushing everything into one PR.

I have read through all of the changes and as far as I can see, it looks good but because of the amount of changes I am not confident that it is bug free :)

Completely understand that. Have really been relying on the CI and tests to catch any issues. Also have been doing some profiling locally, which has helped as well. Plus regex has been my friend (most changes represent a common pattern).

As a side note I agree, we should get rid of all the unused **kwargs in a future PR. It makes a harder to follow variables through the code.

Agreed. I tried to get rid of them in this PR, but it wound up being a little too complicated and touching too much other code. Also as we are still calling the transition_* functions with their Python APIs, we don't get much benefit yet of adding C APIs. Though this may change with future PRs.

[mrocklin]

Actually, I think we should go even further and encapsulate all states of the scheduler and tasks with clearly defined invariants but that is future work :)

Thanks Mads! smile Yeah for sure. I think this opens a lot of doors for us in terms of future optimizations. Just tried to find the right balance where we do enough to get some benefit without pushing everything into one PR.

Yeah, I have the same question. What do we include in SchedulerState and what do we include in Scheduler? I'm totally fine doing this in the future, but some methods, like new_task or transition or remove_worker I expected to be primarily state-focused while methods like stimulus_task_finished or handle_worker I would expect to remain in the Scheduler.

[jakirkham]

Yeah that makes sense. Generally have been thinking anything that relates to the management of the task graph or transformations of it seem like good candidates for SchedulerState. Methods involved in communication or general upkeep/maintenance of the Scheduler seem better there. IOW I think we have the same idea in mind that's just broadly what I've been thinking. Agree that these things can be tackled in future work. Am mainly trying to provide a framework and initial first step here to allow us to dive into that work (possibly from multiple directions 😉).

[jakirkham]

Refactored out a few more minor utility functions used by the transition_*s. AFAICT that's all of them. Please let us know if there's anything else needed here 🙂

[jakirkham]

FWIW here's the profile using Ben's shuffle benchmark with these changes.

[lr4d]

Just wondering, is it necessary that Scheduler inherits from SchedulerState? From a logical level, it seems like it would make more sense if Scheduler had one or mulitple attribute/property provided by SchedulerState

[jakirkham]

After having done the work it's actually not that clear to me that it would be that doable. SchedulerState and Scheduler overlap on a fair number of attributes that they both need. Subclassing makes more sense as they both access these as opposed to letting the Scheduler muck with SchedulerState's internal state, which may actually make it harder to refactor out more functionality later. As the goal was to turn some of the Scheduler's methods into C functions using Cython, this also was easier to do with minimal duplication or disruption via subclassing. Also some of the methods like worker_objective seem to be designed to be overridden in subclasses, which we continue to allow with this structure. Maybe over time as the boundaries between these two become more clearly defined we can revisit, but at least today I'm not sure this is doable.

[mrocklin]

It might not help any, but any thoughts on pulling these out of the transition function into transitions ?

[mrocklin]

It looks like they're a decent chunk of the current time spent

[jakirkham]

It's an interesting suggestion and agree it's worth trying. That said, I feel like this PR has gotten quite long and am worried that has slowed down the reviewing of it. Instead would propose we focus on getting this in. That would make it easier to make smaller more focused PRs on moving communication further out and consolidating messages further as well as allowing us to more effectively Cythonize any remaining functions we still see popping up in the profiles. Thoughts? 🙂

(Should add if we want to do this after the patch release tomorrow that also makes sense to me)

[jakirkham]

This was mostly done in PR ( #4451 ). Handling the remaining ones in PR ( #4526 )

[mrocklin]

OK, we're good to go here. I expect that over time we'll move more things over, but I agree with @jakirkham that what is here now is a good firm base.

[jakirkham]

Thanks Matt! 😄

[jakirkham]

Yeah, I have the same question. What do we include in SchedulerState and what do we include in Scheduler? I'm totally fine doing this in the future, but some methods, like new_task or transition or remove_worker I expected to be primarily state-focused while methods like stimulus_task_finished or handle_worker I would expect to remain in the Scheduler.

FWIW moving new_task over in PR ( #4527 ). This was pretty easy to do.

I think transition should eventually be doable. Would be good to more strongly type the arguments to this function first though.

Looked at remove_worker and some pieces may be amenable to moving, but it is also an async method with some IO bits. So we might need to split it into a helper method that lives in SchedulerState and leave the rest in Scheduler. Will need to think about this more.