Eagerly update aggregate statistics for `TaskPrefix` instead of calculating them on-demand #8681

hendrikmakait · 2024-06-07T07:41:45Z

Tests added / passed
Passes pre-commit run --all-files

github-actions · 2024-06-07T08:26:23Z

Unit Test Results

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

29 files ± 0 29 suites ±0 11h 3m 47s ⏱️ + 1h 19m 38s
4 057 tests - 5 3 953 ✅ + 5 97 💤 - 9 7 ❌ +3
55 883 runs +7 589 53 712 ✅ +7 354 2 163 💤 +249 8 ❌ +3

For more details on these failures, see this check.

Results for commit c84d42d. ± Comparison against base commit 5708bdf.

This pull request removes 14 and adds 9 tests. Note that renamed tests count towards both.

distributed.protocol.tests.test_arrow
distributed.protocol.tests.test_collection
distributed.protocol.tests.test_highlevelgraph
distributed.protocol.tests.test_numpy
distributed.protocol.tests.test_pandas
distributed.shuffle.tests.test_graph
distributed.shuffle.tests.test_merge
distributed.shuffle.tests.test_merge_column_and_index
distributed.shuffle.tests.test_metrics
distributed.shuffle.tests.test_rechunk
…

distributed.diagnostics.tests.test_memray ‑ test_basic_integration_scheduler
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_scheduler_report_args[False]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_scheduler_report_args[report_args0]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_workers[1]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_workers[False]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_workers[True]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_workers_report_args[False]
distributed.diagnostics.tests.test_memray ‑ test_basic_integration_workers_report_args[report_args0]
distributed.tests.test_scheduler ‑ test_task_group_and_prefix_statistics

♻️ This comment has been updated with latest results.

hendrikmakait · 2024-06-07T08:54:33Z

distributed/scheduler.py

-        like ``{"memory": 10, "processing": 3, "released": 4, ...}``
-        """
-        return merge_with(sum, [tg.states for tg in self.groups])
+    @_deprecated(use_instead="groups")  # type: ignore[misc]


The fact that we remove groups from prefixes once all their tasks have been forgotten in

distributed/distributed/scheduler.py

Lines 2031 to 2035 in 5708bdf

if ts.state == "forgotten" and tg.name in self.task_groups:

# Remove TaskGroup if all tasks are in the forgotten state

if all(v == 0 or k == "forgotten" for k, v in tg.states.items()):

ts.prefix.groups.remove(tg)

del self.task_groups[tg.name]

means that the additional filtering applied to active in

distributed/distributed/scheduler.py

Lines 995 to 1001 in 5708bdf

@property

def active(self) -> list[TaskGroup]:

return [

tg

for tg in self.groups

if any(k != "forgotten" and v != 0 for k, v in tg.states.items())

]

is superfluous.

hendrikmakait · 2024-06-07T09:25:37Z

distributed/scheduler.py

+class TaskPrefix(TaskCollection):
+    """Collection tracking all tasks within a prefix
+
+    # FIXME: This comment belongs to the TaskGroup


TODO: I need to adjust this

fjetter

LGTM modulo minor nits

fjetter · 2024-06-07T10:19:04Z

distributed/scheduler.py

+        group.add(self)
+        group.prefix.add(self)


Intuitively, I would expect group.add(self) to be sufficient and not require the caller to also do group.prefix.add(self).

Why can't TaskGroup.add be extended to call TaskPrefix.add?

Same question for transition below

Fair point, this is now aligned.

fjetter · 2024-06-07T10:25:31Z

distributed/scheduler.py

+        self._groups.pop(tg)
+        for state, count in tg.states.items():
+            self.states[state] -= count
+        self.duration -= tg.duration


I guess we've had this problem before but I strongly suspect we'll have to deal with floating point arithmetic precision here.

Historically, we've occasionally encountered negative occupancy issues and maybe this goes back to this diff.

To avoid this issue entirely we may want to track duration as an integer with ms accuracy to avoid this problem. This would require a bit of refactoring and is likely out of scope for this PR.

A quick fix would be to add self.duration = max(self.duration, 0) below this

I think the duration tracking is only loosely-coupled with other durations in the system, so internally tracking this in ms-resolution and exposing the seconds-variant should be pretty quick to do without far-reaching consequences.

2306ae7 stores the duration in microseconds internally. (I think that's the smallest precision we get, milliseconds are too coarse-grained and cause test failures.)

fjetter · 2024-06-07T10:31:27Z

distributed/scheduler.py

+        self.group.transition(self._state, value)
+        self.prefix.transition(self._state, value)
        self._state = value
-        self.prefix.state_counts[value] += 1



Something obvious: I think this is the only truly performance critical section since this is effectively called on every task state transition. The duration and nbytes update is only called whenever a task completes so at least five times more rarely. Since this essentially performs the same operations (plus a couple of additional function indirections/calls) this should perform similarly to the old version

similarly, the nbytes and duration updates should now be twice as expensive but that's a constant. Given that we iterated over everything and performed this computation every 100ms before I assume this is a net positive

fjetter · 2024-06-07T10:34:24Z

distributed/scheduler.py

            if ts.state == "forgotten" and tg.name in self.task_groups:
                # Remove TaskGroup if all tasks are in the forgotten state
                if all(v == 0 or k == "forgotten" for k, v in tg.states.items()):
-                    ts.prefix.groups.remove(tg)
+                    ts.prefix.remove_group(tg)
                    del self.task_groups[tg.name]



I wonder if this cannot be encapsulated in the TaskState.transition method

Well, I guess not entirely. The del self.task_groups[tg.name] would still have to be done here. never mind then

hendrikmakait added 3 commits June 7, 2024 08:56

Improve tests

9ffb2cf

Check groups and active

4a5124e

Rewrite (WIP)

5ee4cac

hendrikmakait added 2 commits June 7, 2024 10:30

Refactor types

2f6fadd

Fix tests

7aa085b

hendrikmakait commented Jun 7, 2024

View reviewed changes

hendrikmakait marked this pull request as ready for review June 7, 2024 09:25

hendrikmakait requested a review from fjetter as a code owner June 7, 2024 09:25

hendrikmakait commented Jun 7, 2024

View reviewed changes

fjetter reviewed Jun 7, 2024

View reviewed changes

hendrikmakait added 2 commits June 7, 2024 13:32

Align

4d5cc4f

store duration in microsecond-granularity

2306ae7

hendrikmakait force-pushed the eagerly-updated-prefix-stats branch from 1114b9e to 2306ae7 Compare June 7, 2024 12:04

hendrikmakait added 3 commits June 7, 2024 14:40

Fix all_durations

98a3a0c

Comment

6bf63b7

Rounding error

c84d42d

hendrikmakait requested a review from fjetter June 7, 2024 16:30

fjetter approved these changes Jun 10, 2024

View reviewed changes

hendrikmakait merged commit 490b696 into dask:main Jun 10, 2024
27 of 34 checks passed

hendrikmakait deleted the eagerly-updated-prefix-stats branch June 10, 2024 09:01

fjetter mentioned this pull request Jun 10, 2024

Warn user if we encounter too many tasks groups on the scheduler #8678

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Eagerly update aggregate statistics for `TaskPrefix` instead of calculating them on-demand #8681

Eagerly update aggregate statistics for `TaskPrefix` instead of calculating them on-demand #8681

hendrikmakait commented Jun 7, 2024

github-actions bot commented Jun 7, 2024 •

edited

Loading

hendrikmakait Jun 7, 2024 •

edited

Loading

hendrikmakait Jun 7, 2024

fjetter left a comment

fjetter Jun 7, 2024

hendrikmakait Jun 7, 2024

fjetter Jun 7, 2024

hendrikmakait Jun 7, 2024

hendrikmakait Jun 7, 2024 •

edited

Loading

fjetter Jun 7, 2024

fjetter Jun 7, 2024

fjetter Jun 7, 2024

fjetter Jun 7, 2024

	if ts.state == "forgotten" and tg.name in self.task_groups:
	# Remove TaskGroup if all tasks are in the forgotten state
	if all(v == 0 or k == "forgotten" for k, v in tg.states.items()):
	ts.prefix.groups.remove(tg)
	del self.task_groups[tg.name]

	@property
	def active(self) -> list[TaskGroup]:
	return [
	tg
	for tg in self.groups
	if any(k != "forgotten" and v != 0 for k, v in tg.states.items())
	]

Eagerly update aggregate statistics for TaskPrefix instead of calculating them on-demand #8681

Eagerly update aggregate statistics for TaskPrefix instead of calculating them on-demand #8681

Conversation

hendrikmakait commented Jun 7, 2024

github-actions bot commented Jun 7, 2024 • edited Loading

Unit Test Results

hendrikmakait Jun 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fjetter 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

hendrikmakait Jun 7, 2024 • edited Loading

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

Eagerly update aggregate statistics for `TaskPrefix` instead of calculating them on-demand #8681

Eagerly update aggregate statistics for `TaskPrefix` instead of calculating them on-demand #8681

github-actions bot commented Jun 7, 2024 •

edited

Loading

hendrikmakait Jun 7, 2024 •

edited

Loading

hendrikmakait Jun 7, 2024 •

edited

Loading