Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Avoid deepcopy when submitting graph #8633

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Avoid deepcopy when submitting graph #8633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
865f080
Dont use topickle when sending graph
fjetter May 2, 2024
a75c993
fix kwarg name
fjetter May 2, 2024
2d5158a
fix update graph test
fjetter May 2, 2024
b3ef5ce
fix
fjetter May 3, 2024
458eb98
Merge branch 'main' into dont_use_topickle_hlg
fjetter May 3, 2024
a990563
use ToPickle
fjetter May 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions distributed/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3150,8 +3150,9 @@ def _graph_to_futures(
futures = {key: Future(key, self, inform=False) for key in keyset}
# Circular import
from distributed.protocol import serialize
from distributed.protocol.serialize import ToPickle
from distributed.protocol.serialize import Pickled, ToPickle

# This is pulled out to have better exception messages
header, frames = serialize(ToPickle(dsk), on_error="raise")

pickled_size = sum(map(nbytes, [header] + frames))
Expand All @@ -3170,8 +3171,7 @@ def _graph_to_futures(
self._send_to_scheduler(
{
"op": "update-graph",
"graph_header": header,
"graph_frames": frames,
"graph": Pickled(header, frames),
"keys": list(keys),
"internal_priority": internal_priority,
"submitting_task": getattr(thread_state, "key", None),
Expand Down
12 changes: 5 additions & 7 deletions distributed/scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@
from distributed.proctitle import setproctitle
from distributed.protocol import deserialize
from distributed.protocol.pickle import dumps, loads
from distributed.protocol.serialize import Serialized, ToPickle, serialize
from distributed.protocol.serialize import Serialize, Serialized, serialize
from distributed.publish import PublishExtension
from distributed.pubsub import PubSubSchedulerExtension
from distributed.queues import QueueExtension
Expand Down Expand Up @@ -3439,7 +3439,7 @@ def _task_to_msg(self, ts: TaskState, duration: float = -1) -> dict[str, Any]:
for dts in ts.dependencies
},
"nbytes": {dts.key: dts.nbytes for dts in ts.dependencies},
"run_spec": ToPickle(ts.run_spec),
"run_spec": Serialize(ts.run_spec),
"resource_restrictions": ts.resource_restrictions,
"actor": ts.actor,
"annotations": ts.annotations or {},
Expand Down Expand Up @@ -4664,8 +4664,7 @@ def _create_taskstate_from_graph(
async def update_graph(
self,
client: str,
graph_header: dict,
graph_frames: list[bytes],
graph: Serialized,
keys: set[Key],
internal_priority: dict[Key, int] | None,
submitting_task: Key | None,
Expand All @@ -4679,8 +4678,7 @@ async def update_graph(
start = time()
try:
try:
graph = deserialize(graph_header, graph_frames).data
del graph_header, graph_frames
graph = deserialize(graph.header, graph.frames)
except Exception as e:
msg = """\
Error during deserialization of the task graph. This frequently
Expand All @@ -4695,7 +4693,7 @@ async def update_graph(
annotations_by_type,
) = await offload(
_materialize_graph,
graph=graph,
graph=graph, # type: ignore[arg-type]
global_annotations=annotations or {},
)
del graph
Expand Down
7 changes: 2 additions & 5 deletions distributed/tests/test_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,8 @@
from distributed.compatibility import LINUX, MACOS, WINDOWS, PeriodicCallback
from distributed.core import ConnectionPool, Status, clean_exception, connect, rpc
from distributed.metrics import time
from distributed.protocol import serialize
from distributed.protocol import Serialize, Serialized, serialize
from distributed.protocol.pickle import dumps, loads
from distributed.protocol.serialize import ToPickle
from distributed.scheduler import KilledWorker, MemoryState, Scheduler, WorkerState
from distributed.utils import TimeoutError, wait_for
from distributed.utils_test import (
Expand Down Expand Up @@ -1431,10 +1430,8 @@ async def test_update_graph_culls(s, a, b):
dependencies={"foo": set()},
)

header, frames = serialize(ToPickle(dsk), on_error="raise")
await s.update_graph(
graph_header=header,
graph_frames=frames,
graph=Serialized(*serialize(Serialize(dsk), on_error="raise")),
keys=["y"],
client="client",
internal_priority={k: 0 for k in "xyz"},
Expand Down
Loading