Use Task class instead of tuple #8797

fjetter · 2024-07-24T11:36:47Z

This is an early version that will close dask/dask#9969

It introduces a new Task class (name is subject to change) and a couple of other related subclasses that should replace the tuple as a representation of runnable tasks.

The benefits of this are outlined in dask/dask#9969 but are primarily focused to reduce overhead during serialization and parsing of results. An important result is also that we can trivially cache functions (and arguments if we wish) to avoid problems like #8767 where users are erroneously providing expensive to pickle functions (which also happens frequently in our own code and/or downstream projects like xarray)

This approach allows us to convert the legacy dsk graph to the new representation with full backwards compatibility. Old graphs can be migrated and new ones written directly using this new representation which will ultimately reduce overhead.

I will follow up with measurements shortly.

Sibling PR in dask dask/dask#11248

fjetter · 2024-07-24T11:37:39Z

distributed/client.py

+                from dask.task_spec import Task

                dsk.update(
                    {
-                        key: (apply, self.func, (tuple, list(args)), kwargs2)
+                        key: Task(key, self.func, args, kwargs2)
+                        # (apply, self.func, (tuple, list(args)), kwargs2)


This isn't necessary but an example of how this migration would look like

fjetter · 2024-07-24T11:37:57Z

distributed/scheduler.py

+        prefix_name = ts.prefix.name
+        count = self.task_prefix_count[prefix_name] - 1
+        tp_count = self.task_prefix_count
+        tp_count_global = self.scheduler._task_prefix_count_global
        if count:
-            self.task_prefix_count[ts.prefix.name] = count
+            tp_count[prefix_name] = count
        else:
-            del self.task_prefix_count[ts.prefix.name]
+            del tp_count[prefix_name]

-        count = self.scheduler._task_prefix_count_global[ts.prefix.name] - 1
+        count = tp_count_global[prefix_name] - 1
        if count:
-            self.scheduler._task_prefix_count_global[ts.prefix.name] = count
+            tp_count_global[prefix_name] = count
        else:
-            del self.scheduler._task_prefix_count_global[ts.prefix.name]
+            del tp_count_global[prefix_name]


this is an unrelated perf fix

I'm curious, how noticeable is this? Apart from that, let's move this to a separate PR to keep this focused on the major change you introduce here.

fjetter · 2024-07-24T11:38:23Z

distributed/scheduler.py

+    dsk = convert_old_style_dsk(dsk)
+    # TODO: This isn't working yet as expected
+    dependencies = dict(DependenciesMapping(dsk))
+
    return dsk, dependencies, annotations_by_type


most/all of this complexity is now either gone entirely or hidden in the class

github-actions · 2024-07-24T11:53:05Z

Unit Test Results

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

25 files ± 0 25 suites ±0 10h 21m 22s ⏱️ + 1m 27s
4 123 tests - 9 4 006 ✅ - 9 110 💤 ±0 7 ❌ +1
47 622 runs - 109 45 521 ✅ - 108 2 087 💤 - 8 14 ❌ +8

For more details on these failures, see this check.

Results for commit 65ed5d5. ± Comparison against base commit 48509b3.

This pull request removes 11 and adds 2 tests. Note that renamed tests count towards both.

distributed.tests.test_client ‑ test_persist_get
distributed.tests.test_client ‑ test_recreate_error_array
distributed.tests.test_client ‑ test_recreate_error_collection
distributed.tests.test_client ‑ test_recreate_error_delayed
distributed.tests.test_client ‑ test_recreate_error_futures
distributed.tests.test_client ‑ test_recreate_task_array
distributed.tests.test_client ‑ test_recreate_task_collection
distributed.tests.test_client ‑ test_recreate_task_delayed
distributed.tests.test_client ‑ test_recreate_task_futures
distributed.tests.test_utils ‑ test_maybe_complex
…

distributed.tests.test_client ‑ test_persist_get[False]
distributed.tests.test_client ‑ test_persist_get[True]

♻️ This comment has been updated with latest results.

jacobtomlinson

There may be implications for some of the dashboard components, the "pew pew pew" plot comes to mind. I see this is still a draft, let me know when it's in a reviewable state and I'll look over the dashboard code to see if anything needs changing there 🙂.

fjetter · 2024-07-24T16:43:53Z

There may be implications for some of the dashboard components, the "pew pew pew" plot comes to mind. I see this is still a draft, let me know when it's in a reviewable state and I'll look over the dashboard code to see if anything needs changing there 🙂.

I'd actually be surprised if that was affected since we don't change the scheduler internal metadata (like dependencies, transfers, where the tasks are executed...). But who knows. I'll probably stumble over 50 small weird things trying to get CI green :)

distributed/scheduler.py

fjetter · 2024-08-23T17:23:36Z

distributed/scheduler.py

+            task_state_created = time()
+            metrics.update(
+                {
+                    "start": start,
+                    "duration_materialization": materialization_done - start,
+                    "duration_ordering": materialization_done - ordering_done,
+                    "duration_state_initialization": ordering_done - task_state_created,
+                    "duration_total": task_state_created - start,
+                }
+            )
+            evt_msg = {
+                "action": "update-graph",
+                "stimulus_id": stimulus_id,
+                "metrics": metrics,
+                "status": "OK",
+            }
+            self.log_event(["all", client, "update-graph"], evt_msg)
+            logger.debug("Task state created. %i new tasks", len(self.tasks) - before)
+        except Exception as e:
+            evt_msg = {
+                "action": "update-graph",
+                "stimulus_id": stimulus_id,
+                "status": "error",
+            }


This is an unrelated change but it shouldn't be too disruptive for the review

Would it be useful to log (partial) metrics on exception? Also, should we add the exception to the log event?

I don't think this is very useful and I don't want to add the exception to this. This is primarily supposed to be a stream of metrics and I don't like adding Exception objects or large strings to it. I also find the logging and proper exception handling mechanism should be sufficient.

If this is a contentious topic I will remove this change from the PR

fjetter · 2024-08-23T17:26:19Z

distributed/scheduler.py

+            logger.debug("Materialization done. Got %i tasks.", len(dsk))
+            dependents = reverse_dict(dependencies)
+            dsk = resolve_aliases(dsk, keys, dependents)
+            dependencies = dict(DependenciesMapping(dsk))
+            logger.debug("Removing aliases. %i left", len(dsk))


This is new... it removes all those

"key": "other-key"

references that we currently schedule as real tasks. Particularly fusion adds these kinds of redirects. I should probably factor this out to a dedicated PR. The impact can be quite substantial in graph size reduction.

fjetter · 2024-08-26T16:41:42Z

distributed/scheduler.py

+    # FIXME: There should be no need to fully materialize and copy this but some
+    # sections in the scheduler are mutating it.
+    dependencies = {k: set(v) for k, v in DependenciesMapping(dsk3).items()}
+    return dsk3, dependencies, annotations_by_type


fjetter

I recommend reviewers to start with the dask/dask PR.

This PR does not include (m)any intentional changes other than adjusting the code to the new spec. There are one or two things that change behavior, I flagged them explicitly

fjetter · 2024-08-26T16:43:29Z

distributed/scheduler.py

+        prefix = ts.prefix
+        duration: float = prefix.duration_average
        if duration >= 0:
            return duration

-        s = self.unknown_durations.get(ts.prefix.name)
+        s = self.unknown_durations.get(prefix.name)
        if s is None:
-            self.unknown_durations[ts.prefix.name] = s = set()
+            self.unknown_durations[prefix.name] = s = set()


this is unrelated (but I won't create another PR for these three lines)

fjetter · 2024-08-26T16:44:23Z

distributed/deploy/tests/test_local.py

+        processes=False,
+        asynchronous=True,
+        scheduler_sync_interval="1ms",
+        dashboard_address=":0",


these dashboard changes are also unrelated. Appologies. If it actually helps I will factor it out but those tests are typically disjoint from actual changes so I hope the review process is not too difficult

This change allows the tests to run in parallel

fjetter · 2024-08-26T16:45:43Z

distributed/scheduler.py

+    # This is removing weird references like "x-foo": "foo" which often make up
+    # a substantial part of the graph
+    # This also performs culling!
+    dsk3 = resolve_aliases(dsk2, keys, dependents)
+


This is an actual change in behavior! This will reduce graphs sizes substantially for graphs that went through linear fusion

distributed/client.py

distributed/tests/test_client.py

hendrikmakait · 2024-08-27T12:42:34Z

distributed/tests/test_client.py

@@ -8402,6 +8237,9 @@ async def test_release_persisted_collection(c, s, a, b):
        await c.compute(arr)


+@pytest.mark.skip(
+    reason="Deadlocks likely related to future serialization and ref counting"


Should we add an issue for this?

that should've been fixed by #8827

I'll remove the skip

distributed/scheduler.py

hendrikmakait · 2024-08-27T12:50:43Z

distributed/scheduler.py

+            metrics.update(
+                {
+                    "start": start,
+                    "duration_materialization": materialization_done - start,


nit: I suggest that we start following Prometheus variable naming conventions here to make our lives easier in the future.

can you suggest the appropriate names? I'm not sure what the correct way is

Suggested change

"duration_materialization": materialization_done - start,

"materialization_duration_seconds": materialization_done - start,

hendrikmakait · 2024-08-27T12:52:09Z

distributed/scheduler.py

+            task_state_created = time()
+            metrics.update(
+                {
+                    "start": start,
+                    "duration_materialization": materialization_done - start,
+                    "duration_ordering": materialization_done - ordering_done,
+                    "duration_state_initialization": ordering_done - task_state_created,
+                    "duration_total": task_state_created - start,
+                }
+            )
+            evt_msg = {
+                "action": "update-graph",
+                "stimulus_id": stimulus_id,
+                "metrics": metrics,
+                "status": "OK",
+            }
+            self.log_event(["all", client, "update-graph"], evt_msg)
+            logger.debug("Task state created. %i new tasks", len(self.tasks) - before)
+        except Exception as e:
+            evt_msg = {
+                "action": "update-graph",
+                "stimulus_id": stimulus_id,
+                "status": "error",
+            }


Would it be useful to log (partial) metrics on exception? Also, should we add the exception to the log event?

hendrikmakait · 2024-08-27T12:53:01Z

distributed/scheduler.py

+                "stimulus_id": stimulus_id,
+                "status": "error",
+            }
+            self.log_event(["all", client, "update-graph"], evt_msg)


Is there a particular reason you prefer a dedicated update-graph topic instead of something like a scheduler topic?

No particular reason. I used this in a test but I will change it. I will also drop the all topic (feels a bit like an anti pattern)

Note that dropping the all pattern is a user-facing breaking change. That being said, I'm all for redesigning our topics, etc., this might just require some changes for downstream users, e.g., for Coiled.

Then again, we've also renamed the action, so these changes are breaking already.

Yeah, I think this is breaking and I'm fine with it. This is overall a pretty burried feature and I doubt (m)any users will notice.

distributed/worker.py

fjetter · 2024-10-15T10:00:44Z

needs dask/dask#11429

fjetter · 2024-10-15T13:48:21Z

Also this dask/dask#11431

…buted into support_task_spec

fjetter · 2024-10-16T13:22:54Z

The mindeps builds are sad. Everything else seems unrelated

fjetter · 2024-10-16T17:00:19Z

distributed/scheduler.py

+        seen: set[Key] = set()
+        sadd = seen.add
+        for k in list(keys):
+            work = {k}
+            wpop = work.pop
+            wupdate = work.update
+            while work:
+                d = wpop()
+                if d in seen:
+                    continue
+                sadd(d)
+                if d not in dsk:
+                    if d not in self.tasks:
+                        lost_keys.add(d)
+                        lost_keys.add(k)
+                        logger.info("User asked for computation on lost data, %s", k)
+                        dependencies.pop(d, None)
+                        keys.discard(k)
+                    continue
+                wupdate(dsk[d].dependencies)
+        return lost_keys


I rewrote this section. I had issues with it and just barely understand the old code (and ran into multiple bugs in the recent past)

hendrikmakait

The code generally looks good to me. CI test results are confusing, it seems like they're out of sync with the actual test jobs?

I've added a bunch of suggestions for Prometheus-style metric naming within update_graph.

continuous_integration/environment-mindeps.yaml

distributed/shuffle/_rechunk.py

distributed/tests/test_utils_comm.py

distributed/scheduler.py

hendrikmakait · 2024-10-18T07:49:02Z

distributed/scheduler.py

+            "new-tasks": len(new_tasks),
+            "key-collisions": colliding_task_count,


Prometheus naming prefers underscores over hyphens

Suggested change

"new-tasks": len(new_tasks),

"key-collisions": colliding_task_count,

"new_tasks": len(new_tasks),

"key_collisions": colliding_task_count,

hendrikmakait · 2024-10-18T07:49:53Z

distributed/scheduler.py

+            metrics.update(
+                {
+                    "start": start,
+                    "duration_materialization": materialization_done - start,


Suggested change

"duration_materialization": materialization_done - start,

"materialization_duration_seconds": materialization_done - start,

hendrikmakait · 2024-10-18T07:52:52Z

distributed/scheduler.py

+                    "duration_ordering": materialization_done - ordering_done,
+                    "duration_state_initialization": ordering_done - task_state_created,
+                    "duration_total": task_state_created - start,


Prometheus naming convention (roughly): (_total if accumulating)

Suggested change

"duration_ordering": materialization_done - ordering_done,

"duration_state_initialization": ordering_done - task_state_created,

"duration_total": task_state_created - start,

"ordering_duration_seconds": materialization_done - ordering_done,

"state_initialization_duration_seconds": ordering_done - task_state_created,

"duration_seconds": task_state_created - start,

(I don't have a great suggestion for the e2e duration)

hendrikmakait · 2024-10-18T07:55:06Z

distributed/scheduler.py

+            task_state_created = time()
+            metrics.update(
+                {
+                    "start": start,


Suggested change

"start": start,

"start_timestamp_seconds": start,

hendrikmakait · 2024-10-18T08:01:59Z

Update: test_merge failures seem systemic (and thus related).

fjetter · 2024-10-18T09:01:22Z

Update: test_merge failures seem systemic (and thus related).

I've been trying to reproduce but without any luck so far

…buted into support_task_spec

fjetter · 2024-10-18T09:52:08Z

I think the test_merge failures are actually unrelated. The exception is

  
        def validate_data(self, data: pd.DataFrame) -> None:
>           if set(data.columns) != set(self.meta.columns):
E           AttributeError: 'tuple' object has no attribute 'columns'

which indicates that data is likely a key instead of the data... I've seen this before and I think this is a dask-expr problem. I added another verification step here to confirm this. I'll keep digging

fjetter · 2024-10-18T10:28:38Z

yeah, so the exception is pretty much what I expected

            if not isinstance(data, pd.DataFrame):
>               raise TypeError(f"Expected {data=} to be a DataFrame, got {type(data)}.")
E               TypeError: Expected data=('assign-3d7cfa7cea412465799bea6cfac1b512', 1) to be a DataFrame, got <class 'tuple'>.

fjetter · 2024-10-18T11:14:56Z

Ah, this is the build with dask-expr disabled. Now I can reproduce!

fjetter · 2024-10-18T11:20:43Z

distributed/shuffle/_shuffle.py

                shuffle_transfer,
-                (self.name_input, i),
+                TaskRef((self.name_input, i)),


The test_merge tests w/ dask-expr enabled never take this code path. That's interesting but not incredibly surprising.

fjetter · 2024-10-18T13:18:17Z

distributed/shuffle/tests/test_graph.py::test_multiple_linear failure is also related. also a legacy-only problem

fjetter · 2024-10-18T15:31:14Z

dask/dask#11445 is hopefully the last one

hendrikmakait · 2024-10-22T15:32:33Z

Is there anything left to do here?

fjetter mentioned this pull request Jul 24, 2024

Add a Task class to replace tuples for task specification dask/dask#11248

Merged

fjetter commented Jul 24, 2024

View reviewed changes

jacobtomlinson reviewed Jul 24, 2024

View reviewed changes

fjetter force-pushed the support_task_spec branch from 99a2ec5 to bb0324a Compare July 29, 2024 13:23

hendrikmakait self-requested a review July 30, 2024 14:52

hendrikmakait reviewed Jul 30, 2024

View reviewed changes

distributed/scheduler.py Show resolved Hide resolved

fjetter force-pushed the support_task_spec branch 3 times, most recently from b01c705 to bb7d38e Compare August 7, 2024 12:56

This was referenced Aug 7, 2024

Avoid excessive attribute access overhead for remove_from_task_prefix_count #8821

Merged

Avoid key validation if validation is disabled #8822

Merged

fjetter force-pushed the support_task_spec branch 2 times, most recently from 507fcb6 to 45b42a5 Compare August 9, 2024 14:05

fjetter force-pushed the support_task_spec branch 3 times, most recently from eb30895 to 2ecaa5c Compare August 23, 2024 16:51

fjetter commented Aug 23, 2024

View reviewed changes

fjetter marked this pull request as ready for review August 23, 2024 17:31

fjetter force-pushed the support_task_spec branch from 919ff4f to 9974bb0 Compare August 26, 2024 11:56

This was referenced Aug 26, 2024

Full support for task spec in dask.order dask/dask#11347

Merged

Dask order uses task_spec #8842

Open

hendrikmakait self-requested a review August 26, 2024 12:50

fjetter commented Aug 26, 2024

View reviewed changes

hendrikmakait reviewed Aug 27, 2024

View reviewed changes

fjetter added 3 commits October 15, 2024 12:05

Merge branch 'main' into support_task_spec

55cf70b

fix typing

4f6f700

fix

eba5490

fjetter added 2 commits October 15, 2024 16:38

Merge branch 'main' into support_task_spec

c1bbbfb

Merge branch 'support_task_spec' of https://github.com/fjetter/distri…

380e33f

…buted into support_task_spec

fjetter mentioned this pull request Oct 16, 2024

Release 2024.10.0 dask/community#401

Closed

4 tasks

fjetter commented Oct 16, 2024

View reviewed changes

hendrikmakait approved these changes Oct 18, 2024

View reviewed changes

fjetter added 5 commits October 18, 2024 11:01

Review

8d6fd39

Merge branch 'main' into support_task_spec

e22e65a

fix key collisions metrics name

f546d00

Merge branch 'support_task_spec' of https://github.com/fjetter/distri…

bb071db

…buted into support_task_spec

Add type validation exception

dd675f0

fix p2p shuffle

8349918

fjetter commented Oct 18, 2024

View reviewed changes

remove whitespace

83c9f3f

use is_dataframe_like

65ed5d5

fjetter merged commit 928d770 into dask:main Oct 24, 2024
25 of 32 checks passed

fjetter deleted the support_task_spec branch October 24, 2024 13:53

	"duration_materialization": materialization_done - start,
	"materialization_duration_seconds": materialization_done - start,

		"new-tasks": len(new_tasks),
		"key-collisions": colliding_task_count,

Use Task class instead of tuple #8797

Use Task class instead of tuple #8797

Conversation

fjetter commented Jul 24, 2024 • edited Loading

fjetter Jul 24, 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

github-actions bot commented Jul 24, 2024 • edited Loading

Unit Test Results

jacobtomlinson left a comment

Choose a reason for hiding this comment

fjetter commented Jul 24, 2024

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

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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fjetter commented Oct 15, 2024

fjetter commented Oct 15, 2024

fjetter commented Oct 16, 2024

Choose a reason for hiding this comment

hendrikmakait 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 commented Oct 18, 2024

fjetter commented Oct 18, 2024

fjetter commented Oct 18, 2024

fjetter commented Oct 18, 2024

fjetter commented Oct 18, 2024

Choose a reason for hiding this comment

fjetter commented Oct 18, 2024

fjetter commented Oct 18, 2024

hendrikmakait commented Oct 22, 2024

fjetter commented Jul 24, 2024 •

edited

Loading

fjetter Jul 24, 2024 •

edited

Loading

github-actions bot commented Jul 24, 2024 •

edited

Loading