Annotate `Task*` objects for Cythonization #4302

jakirkham · 2020-12-03T04:09:31Z

Analogous to PR ( #4290 ) ( #4294 ) except this is annotating TaskState, TaskPrefix, and TaskGroup. Plus all usages thereof. This is a bit longer simply because of how many attributes TaskState has and how frequently it is used. That said, it follows the same pattern as was seen with the other two and so shouldn't be too surprising.

mrocklin · 2020-12-03T04:12:05Z

I'm excited about seeing this one in. It's probably too early, but I'd be curious to see how this changes profiles, if at all.

…

On Wed, Dec 2, 2020, 8:09 PM jakirkham ***@***.***> wrote: Requires PR ( #4294 <#4294> ) Analogous to PR ( #4290 <#4290> ) ( #4294 <#4294> ) except this is annotating TaskState and usages thereof. This is a bit longer simply because of how many attributes TaskState has and how frequently it is used. That said, it follows the same pattern as was seen with the other two and so shouldn't be too surprising. As TaskState and WorkerState overlap a lot, this just builds on top of PR ( #4294 <#4294> ) to avoid conflicts. Will rebase this after PR ( #4294 <#4294> ) is in. Though PR ( #4294 <#4294> ) should be reviewed and merged before worrying about this one 😉 ------------------------------ You can view, comment on, or merge this pull request online at: #4302 Commit Summary - Use `ws` variable name for `WorkerState` objects - Name `WorkerState` variable distinctly in closure - Assign selected `WorkerState` to variable - Annotate `WorkerState` for Cythonization - Annotate all `WorkerState` variables - Add Python-level `property`s for attributes - Add some `property.setter`s - Create `list` from generator - Run `black` - Relax `_address` to `object` - Relax `_name` to `object` - Use `tsp` variable name for `TaskStreamPlugin`s - Use `dts` for iterated `TaskState` variables - Create `list` from generator - Use `-1` as `TaskState.nbytes` default - Assign `TaskState` instances to variables - Annotate `TaskState` for Cythonization - Annotate all `TaskState` variables - Use closure to access `TaskState.priority` - Add `_` before all `TaskState` attributes - Use `_` prefixed `TaskState` attributes throughout - Add Python-level `property`s for attributes - Add some `property.setter`s - Drop recently added `TaskState.priority` closures File Changes - *M* distributed/scheduler.py <https://github.com/dask/distributed/pull/4302/files#diff-bbcf2e505bf2f9dd0dc25de4582115ee4ed4a6e80997affc7b22122912cc6591> (1934) Patch Links: - https://github.com/dask/distributed/pull/4302.patch - https://github.com/dask/distributed/pull/4302.diff — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#4302>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTARFXCAFFSSC37EYVTSS4FQTANCNFSM4ULNBWQA> .

jakirkham · 2020-12-03T06:50:25Z

Both PR ( #4294 ) and this help based on my own local profiling and looking at the call graphs. Though the other thing they do is unlock further optimizations in the Scheduler transition methods. Since all of those methods are just working with these 3 objects ClientState, WorkerState, and TaskState. I think once we get these in, we will want to take another look at line profiling of the transition methods and see what sticks out. There are different ways we can structure the code and make additional annotations so that Cython can optimize it more effectively.

Even just the ClientState and check_idle_saturated changes helped as evidenced by the nightly profile ( quasiben/dask-scheduler-performance#39 ).

Anyways if you have time to review PR ( #4294 ) tomorrow that will be really helpful. This PR is 95% there, but there is one or two more things that it may still need. So will be working on them as well.

mrocklin

Some minor comments here. In general all of this looks straightforward.

I'm curious about a couple of the types, but those can wait as well.

Also beware, you're likely to get a few conflicts when we merge in the annotations work #4279

distributed/scheduler.py

mrocklin · 2020-12-03T15:23:22Z

distributed/scheduler.py

+    _key: str
+    _hash: Py_hash_t
+    _prefix: object
+    _run_spec: object


bytes maybe

Interesting. This came from the docstring above. Later in the code it was implied this could be a dict, but might not be. Seems like it is not very clear what it may be so perhaps it is best to leave as object for now.

mrocklin · 2020-12-03T15:23:52Z

distributed/scheduler.py

+    _prefix: object
+    _run_spec: object
+    _priority: tuple
+    _state: str


Should we make this into an enum? Would Cython prefer that?

Cython can't do much with a Python Enum. There are enums in Cython that it does expose to the Python layer optionally. Though these are not available in pure Python mode unfortunately ( cython/cython#3923 ).

For now I think it is ok to just type this as a str as part of our broad effort to type things. We can then revisit specifics once everything is in and we have had an opportunity to profile more closely.

Sounds good. We check and change state pretty often. I wouldn't be surprised if this has some effect in the future.

Are pure Python enums something that would make sense to request upstream, or is this likely to be hard to achieve with cython?

Yeah it may. I think line profiling after this lands makes sense and should hopefully guide us to which of several things we should further optimize. It could be Status or it could be something else.

Seems like a reasonable request to me (though maybe I'm biased as I already made that request 😄). Have no idea if it will be easy or not. Certainly hope it is doable.

In any event, I think at a first pass having str here is fine. Am more interested in broadly typing things first and worrying about more tuning in subsequent PRs after more profiling.

One other thing I forgot to note, Cython does intern any str literals it sees by default as part of module initialization (IOW at import time), which it hangs onto globally (so they don't go out of scope). That way, when we do something like ts._state = "running", Cython will already have interned "running" and assigned the interned str variable to ts._state. Further when we later do something like assert ts._state == "running", this "running" will also refer to the same interned str. As a result comparisons of strs defined with literals should benefit from the intern speedup.

Though there is a caveat. This does not apply for any dynamically generated strs. So "My name is %s" % s will construct a str at runtime that is not interned.

distributed/scheduler.py

jakirkham · 2020-12-03T15:47:20Z

Since this just includes the WorkerState PR, maybe we should work on that one first? 😉 Unless we’d rather just work with both of them in this PR? 🙂

mrocklin · 2020-12-03T15:48:54Z

I don't have a preference where the work is done. It'd be nice to merge one and then the other for git history. It was easier for me to comment here just because everything was here.

jakirkham · 2020-12-03T16:46:42Z

Understood. If we'd like to merge two separate PRs, would suggest to have WorkerState comments in that PR at least. Otherwise I think once that PR gets merged the comments here may wind up being artificially "resolved" since that code is no longer here, which may make hard for us to ensure we have addressed all of them.

That said, if we find it simply more convenient to work here, it may make more sense to close out that PR and do all the work here. It probably gets harder to disambiguate where changes apply if we start pushing both TaskState and WorkerState changes here.

jakirkham · 2020-12-03T20:00:10Z

Copied the WorkerState comments to PR ( #4294 ). Let's follow up on those over there 🙂

jakirkham · 2020-12-04T16:07:29Z

Rebased on master to get rid of the WorkerState bits.

jakirkham · 2020-12-04T17:07:40Z

distributed/scheduler.py

+    _resource_restrictions: dict
+    _loose_restrictions: bool
+    _metadata: dict
+    _annotations: dict


Also included changes related to annotations from PR ( #4279 ). AFAICT this is just a dict, but please let me know if I'm missing something.

jakirkham · 2020-12-04T20:31:39Z

Do we want to handle Cythonization of TaskPrefix and TaskGroup in this PR? They are pretty small classes and don't see as much use as TaskState. So think it should be straightforward to do. That said, don't want to overload this PR since it already has a lot of changes. WDYT?

Also would be good to get your thoughts on run_spec as the typing seems unclear ( #4302 (comment) ). Am also ok just leaving it as object if we are not sure. We can always follow up later if needed.

Thanks for the feedback thus far 🙂

mrocklin · 2020-12-04T21:59:16Z

No preference on TaskPrefix/TaskGroup. I think that you probably care more about git history than I do.

It would be good to make run_spec more consistent, but it's rarely accessed, so I don't think that this is a priority.

jakirkham · 2020-12-04T22:04:37Z

Was more concerned about how that affects the review load for you. No strong feelings on commit history from me.

Ok let's skip run_spec for now and we can revisit after this lands if it becomes important.

mrocklin · 2020-12-04T22:06:43Z

Review of this stuff is pretty simple. There isn't any tricky logic to go over. I'm generally happy.

jakirkham · 2020-12-04T23:19:42Z

Found I was having issues getting Cython to use the type definitions correctly when quoted. So am proposing just swapping the ordering of Task* class definitions around to make things easier. ( #4302 )

jakirkham · 2020-12-05T01:43:25Z

Need to flip the order of TaskPrefix and TaskGroup. Didn't realize the latter has an attribute that holds the former. Fixing with PR ( #4319 ).

Now that we have properties accessible from `TaskState`, drop the closures we added previously to access the typed values internally. This should be equivalently performant and cut out a little bit of boilerplate.

Make sure to assign to the `TaskPrefix` variable, `tp`, first before assigning to the `dict`. This should avoid the admittedly likely low overhead of looking up the result in the dictionary when we already have the value available.

Saves us need to fetch this twice. Also makes the code a bit more readable. Finally may allow Cython optimizations on the variable later.

Instead of using `None` for `TaskPrefix.duration_average`, set it `-1`. This works better when typing `TaskPrefix.duration_average` as it can always be floating point. This also works logically with this value as it can't actually be negative unless it wasn't defined. Rework the logic around this variable to ensure it is positive semi-definite.

This allows Cython to perform C-level optimizations on these variables and usages thereof.

Make sure to assign to the `TaskGroup` variable, `tg`, first before assigning to the `dict`. This should avoid the admittedly likely low overhead of looking up the result in the dictionary when we already have the value available.

This ensures Cython still uses `TaskGroup` to annotate the variable iterated over. Otherwise it constructs a generator with its own scope where this is ignored.

jakirkham · 2020-12-07T18:27:32Z

Ok this should be good to go. Please let me know if anything else is needed 🙂

quasiben · 2020-12-07T20:20:48Z

I looked over the scheduler and seems like you are following the same techniques you laid out earlier -- _var with type and property decorators for Python access of the variables. Thanks @jakirkham !!

Merging in now

jakirkham · 2020-12-07T20:30:27Z

Thanks Ben! 😄

If anything else comes up, happy to follow up in a new issue/PR as appropriate.

mrocklin reviewed Dec 3, 2020

View reviewed changes

jakirkham mentioned this pull request Dec 3, 2020

Annotate WorkerState for Cythonization #4294

Merged

jakirkham mentioned this pull request Dec 3, 2020

Coerce new TaskState.nbytes value to int #4311

Merged

jakirkham force-pushed the cy_ts branch 3 times, most recently from f89eaaa to aaeeb9c Compare December 4, 2020 16:06

jakirkham commented Dec 4, 2020

View reviewed changes

jakirkham force-pushed the cy_ts branch 2 times, most recently from 6198ee2 to bdafb41 Compare December 4, 2020 20:26

jakirkham mentioned this pull request Dec 4, 2020

Move TaskGroup & TaskPrefix before TaskState #4318

Merged

jakirkham force-pushed the cy_ts branch 5 times, most recently from 842be89 to 182ece9 Compare December 5, 2020 00:59

jakirkham mentioned this pull request Dec 5, 2020

Swap order of TaskGroup and TaskPrefix #4319

Merged

jakirkham force-pushed the cy_ts branch from 182ece9 to 3856340 Compare December 5, 2020 02:03

jakirkham added 19 commits December 7, 2020 09:33

Drop recently added TaskState.priority closures

af9461f

Now that we have properties accessible from `TaskState`, drop the closures we added previously to access the typed values internally. This should be equivalently performant and cut out a little bit of boilerplate.

Swap assignment order with TaskPrefix

e834d88

Make sure to assign to the `TaskPrefix` variable, `tp`, first before assigning to the `dict`. This should avoid the admittedly likely low overhead of looking up the result in the dictionary when we already have the value available.

Assign dask.config query to a variable

be1957c

Saves us need to fetch this twice. Also makes the code a bit more readable. Finally may allow Cython optimizations on the variable later.

Annotate TaskPrefix for Cythonization

3b4b76e

Annotate TaskPrefix constructor as well

cabe00d

Annotate all TaskPrefix variables

47d77b1

Add _ before all TaskPrefix attributes

9ff0dbb

Add Python-level propertys for attributes

9593fb0

Use Py_ssize_t in set_nbytes

5a76c8f

This allows Cython to perform C-level optimizations on these variables and usages thereof.

Use tg for TaskGroup variables

5af80ca

Swap assignment order with TaskGroup

9802c86

Make sure to assign to the `TaskGroup` variable, `tg`, first before assigning to the `dict`. This should avoid the admittedly likely low overhead of looking up the result in the dictionary when we already have the value available.

Create list from generator

53fd21e

This ensures Cython still uses `TaskGroup` to annotate the variable iterated over. Otherwise it constructs a generator with its own scope where this is ignored.

Annotate TaskGroup for Cythonization

1414ff5

Annotate TaskGroup constructor as well

9e2aa4f

Set TaskGroup.prefix to None initially

cc0dbb1

Annotate all TaskGroup variables

1f50eb2

Add _ before all TaskGroup attributes

9336f2a

Add Python-level propertys for attributes

4aa4cbe

jakirkham force-pushed the cy_ts branch from fc4a145 to 4aa4cbe Compare December 7, 2020 17:33

jakirkham changed the title ~~[WIP] Annotate Task* objects for Cythonization~~ Annotate Task* objects for Cythonization Dec 7, 2020

jakirkham marked this pull request as ready for review December 7, 2020 18:26

jakirkham requested a review from mrocklin December 7, 2020 18:54

quasiben merged commit 0d4a499 into dask:master Dec 7, 2020

jakirkham deleted the cy_ts branch December 7, 2020 20:29

jakirkham mentioned this pull request Dec 9, 2020

line_profiler results on 4 workers (w/o stealing) over 20 iterations quasiben/dask-scheduler-performance#20

Open

jakirkham mentioned this pull request Jan 27, 2022

Dask Futures performance as the number of tasks increases #5715

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Annotate `Task*` objects for Cythonization #4302

Annotate `Task*` objects for Cythonization #4302

jakirkham commented Dec 3, 2020 •

edited

Loading

mrocklin commented Dec 3, 2020 via email

jakirkham commented Dec 3, 2020

mrocklin left a comment

mrocklin Dec 3, 2020

jakirkham Dec 4, 2020

mrocklin Dec 3, 2020

jakirkham Dec 4, 2020

mrocklin Dec 4, 2020

jakirkham Dec 4, 2020

jakirkham Dec 10, 2020

jakirkham commented Dec 3, 2020

mrocklin commented Dec 3, 2020

jakirkham commented Dec 3, 2020

jakirkham commented Dec 3, 2020

jakirkham commented Dec 4, 2020

jakirkham Dec 4, 2020

jakirkham commented Dec 4, 2020

mrocklin commented Dec 4, 2020

jakirkham commented Dec 4, 2020

mrocklin commented Dec 4, 2020

jakirkham commented Dec 4, 2020

jakirkham commented Dec 5, 2020

jakirkham commented Dec 7, 2020

quasiben commented Dec 7, 2020

jakirkham commented Dec 7, 2020

Annotate Task* objects for Cythonization #4302

Annotate Task* objects for Cythonization #4302

Conversation

jakirkham commented Dec 3, 2020 • edited Loading

mrocklin commented Dec 3, 2020 via email

jakirkham commented Dec 3, 2020

mrocklin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakirkham commented Dec 3, 2020

mrocklin commented Dec 3, 2020

jakirkham commented Dec 3, 2020

jakirkham commented Dec 3, 2020

jakirkham commented Dec 4, 2020

Choose a reason for hiding this comment

jakirkham commented Dec 4, 2020

mrocklin commented Dec 4, 2020

jakirkham commented Dec 4, 2020

mrocklin commented Dec 4, 2020

jakirkham commented Dec 4, 2020

jakirkham commented Dec 5, 2020

jakirkham commented Dec 7, 2020

quasiben commented Dec 7, 2020

jakirkham commented Dec 7, 2020

Annotate `Task*` objects for Cythonization #4302

Annotate `Task*` objects for Cythonization #4302

jakirkham commented Dec 3, 2020 •

edited

Loading