Reduce memory usage of scheduler process - Optimize scheduler.py::TaskState class

[milesgranger]

Part of #7998

The first approach in ef4b7dd moved things into descriptors, which would only assign a value when being accessed. This results in ~34% reduction in memory when compared to main. Downside was there is some boilerplace for each attribute requiring a lazy creation (albeit this could be generalized by prefixing each with __lazy_ for example then iterating on those annotations after initialization..)

Second approach, the current direction of this PR, uses optionals for more attributes in this class. Primarily those using set/dict. This results in ~42% reduction in memory when compared to main. Downside is much more changed code, checking for None everywhere and the tests are currently wildly broken. 😅

Snippet I'm using for this comparison to main:

from dask.datasets import timeseries
from dask.distributed import Client
from distributed.diagnostics.memray import memray_scheduler

def main():
    ddf = timeseries("2020", "2021", partition_freq='2h')
    ddf2 = ddf.shuffle(on="x", shuffle='tasks')

    print(len(ddf2.__dask_graph__()))

    with Client() as client:
        with memray_scheduler():
            fut = client.compute(ddf2.size)
            import time
            time.sleep(90)
        fut.cancel()

if __name__ == '__main__':
    main()

cc @fjetter

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       27 files  ±0         27 suites  ±0   13h 41m 1s ⏱️ - 1h 13m 13s
  3 926 tests ±0    3 803 ✔️ +4     117 💤 ±0  6 ❌  - 4 
49 391 runs  ±0  46 970 ✔️ +5  2 414 💤 +1  7 ❌  - 6 

For more details on these failures, see this check.

Results for commit 45447a6. ± Comparison against base commit 0dc9e88.

♻️ This comment has been updated with latest results.

[milesgranger]

Okay, @fjetter I think this is ready for a proper review then. Failing tests don't seem related, (ran most/all of those failing here locally fine as well).

[fjetter]

I think summarizing my review

• If possible we should use assert foo is not None instead of if foo else () since it helps readability a lot. The scheduler is a messy, complex system and knowing for a fact that an attribute exists helps
• there are a bunch of places where we're clearing existing dicts/sets. In the spirit of only keeping what's necessary, we can instead just reset them to None
• There are a couple of very common sets that are used for every task and I think we should just keep them non-optional. This concerns primarily dependencies and dependents. They are all but guaranteed to be populated so we can initialize them right away and save ourselves a bit of complexity
• There are a couple of sets that are going to be populated for every single task but only for a certain amount of time. This includes who_has, waiters, waiting_on. We will only save ourselves memory by delaying initialization if we ensure that the sets are actually released once those collections are no longer required (which is typically true once a task is in memory). See the above comment about clear

[fjetter]

I changed add_replica to be idempotent so this validation is no longer valid

[fjetter]

test failures are all unrelated

[fjetter]

test failures are all unrelated

kind of sad considering all the failures...